BMSG-GAN

[MSG-GAN] 的 PyTorch 实现。

**请注意，这不是 MSG-GAN 研究论文的代码仓库。请前往 msg-stylegan-tf 仓库，获取 MSG-GAN 论文的官方代码和训练好的模型。

SageMaker

现在支持在 AWS SageMaker 上进行训练。请参阅 https://docs.aws.amazon.com/sagemaker/latest/dg/pytorch.html

MSG-GAN：用于稳定图像合成的多尺度梯度生成对抗网络

摘要:
尽管生成对抗网络（GAN）在图像合成任务中取得了巨大成功，但它们却以难以使用而闻名，部分原因在于训练过程中的不稳定性。造成这种不稳定的常见原因之一是，从判别器传递到生成器的梯度会因训练过程中的学习不平衡而迅速变得缺乏信息量。在本工作中，我们提出了多尺度梯度生成对抗网络（MSG-GAN），这是一种简单而有效的技术，能够解决这一问题，它允许梯度在多个尺度上从判别器流向生成器。该技术提供了一种稳定的方法来生成同步的多尺度图像。我们提出了一种非常直观的数学 MSG-GAN 框架实现，其中在判别器的计算中使用了拼接操作。我们通过在 CIFAR10 和 Oxford102 花卉数据集上的实验，实证验证了 MSG-GAN 方法的效果，并将其与其他进行多尺度图像合成的相关技术进行了比较。此外，我们还提供了在 CelebA-HQ 数据集中合成 1024×1024 高分辨率图像的详细实验结果。

这是一段关于 MSG-GAN 的解释性训练延时视频/GIF。较高分辨率的层最初显示为纯色块，但最终（很快）训练会渗透到所有层，随后它们协同工作以生成更好的样本。请观察训练的前几秒，可以看到类似人脸的块状结构按照从低分辨率到高分辨率的顺序依次出现。

多尺度梯度架构

上图描述了用于生成同步多尺度图像的 MSG-GAN 架构。我们的方法基于 proGAN 提出的架构，但不同于逐步增长的训练方案，它在生成器的中间层与判别器的中间层之间建立了连接。输入到判别器的多尺度图像会被转换为空间体积，并与卷积层主路径上得到的相应激活体积进行拼接。

在判别过程中，真实图像的适当下采样版本会按图示（从上往下）输入到判别器的对应层。

上图解释了在训练过程中，MSG-GAN 中的所有层首先在颜色上同步，随后在各个尺度上逐步提升生成图像的质量。最终，所有层（尺度）的图像亮度都会实现同步。

运行代码

请注意，在所有实验中，为了获得最佳效果，请为生成器和判别器都使用 learning_rate=0.003 的学习率。该模型非常稳健，即使在不同的学习率设置下，也能很快收敛到非常接近的 FID 或 IS 值。 请使用 relativistic-hinge 作为损失函数（默认设置）进行训练。

通过运行 sourcecode/ 目录下的 train.py 脚本来开始训练。以下参数可供您根据需要调整：

-h, --help            显示帮助信息并退出
  --generator_file GENERATOR_FILE
                        生成器的预训练权重文件
  --generator_optim_file GENERATOR_OPTIM_FILE
                        生成器优化器的保存状态
  --shadow_generator_file SHADOW_GENERATOR_FILE
                        影子生成器的预训练权重文件
  --discriminator_file DISCRIMINATOR_FILE
                        判别器的预训练权重文件
  --discriminator_optim_file DISCRIMINATOR_OPTIM_FILE
                        判别器优化器的保存状态
  --images_dir IMAGES_DIR
                        图像数据目录的路径
  --folder_distributed FOLDER_DISTRIBUTED
                        图像数据目录是否包含子文件夹
  --flip_augment FLIP_AUGMENT
                        是否在训练过程中随机翻转图像
  --sample_dir SAMPLE_DIR
                        生成样本的输出目录路径
  --model_dir MODEL_DIR
                        模型保存目录路径
  --loss_function LOSS_FUNCTION
                        使用的损失函数：standard-gan、wgan-gp、
                        lsgan、lsgan-sigmoid、hinge、relativistic-hinge
  --depth DEPTH         GAN 的深度
  --latent_size LATENT_SIZE
                        生成器的潜在空间大小
  --batch_size BATCH_SIZE
                        训练时的批量大小
  --start START         开始的起始 epoch 数
  --num_epochs NUM_EPOCHS
                        训练的总 epoch 数
  --feedback_factor FEEDBACK_FACTOR
                        每个 epoch 生成的日志数量
  --num_samples NUM_SAMPLES
                        用于生成网格图的样本数量
                        最好是平方数
  --checkpoint_factor CHECKPOINT_FACTOR
                        每 n 个 epoch 保存一次模型
  --g_lr G_LR           生成器的学习率
  --d_lr D_LR           判别器的学习率
  --adam_beta1 ADAM_BETA1
                        Adam 优化器的 beta_1 值
  --adam_beta2 ADAM_BETA2
                        Adam 优化器的 beta_2 值
  --use_eql USE_EQL     是否使用等化学习率
  --use_ema USE_EMA     是否使用指数移动平均
  --ema_decay EMA_DECAY
                        移动平均的衰减系数
  --data_percentage DATA_PERCENTAGE
                        使用的数据百分比
  --num_workers NUM_WORKERS
                        并行读取文件的工作线程数

示例训练命令

若要训练分辨率为 256 x 256 的网络，可使用以下参数：

$ python train.py --depth=7 \ 
                  --latent_size=512 \
                  --images_dir=<图像数据路径> \
                  --sample_dir=samples/exp_1 \
                  --model_dir=models/exp_1

请根据实际情况设置 batch_size、feedback_factor 和 checkpoint_factor。 我们在实验中使用了 DGX-1 机器上的两块 Tesla V100 GPU。

不同数据集上的生成样本

:star: [NEW] :star: CelebA HQ [1024 x 1024] (3万张数据)

:star: [NEW] :star: 牛津花卉数据集（改进版）[256 x 256] (8千张数据)

CelebA HQ [256 x 256] (3万张数据)

LSUN 卧室数据集 [128 x 128] (300万张数据)

CelebA 数据集 [128 x 128] (20万张数据)

全分辨率同步生成样本

Cifar-10 数据集 [32 x 32] (5万张数据)

牛津-102 花卉数据集 [256 x 256] (8千张数据)

引用我们的工作

@article{karnewar2019msg,
  title={MSG-GAN: 多尺度梯度 GAN 用于稳定图像生成},
  author={Karnewar, Animesh 和 Wang, Oliver 以及 Iyengar, Raghu Sesha},
  journal={arXiv 预印本 arXiv:1903.06048},
  year={2019}
}

其他贡献者 :smile:

卡通数据集 [128 x 128] (1万张数据) 由 @huangzh13 提供

感谢

如果您使用此架构在其他数据集上进行了训练，请随时在此处提交 PR。

此致敬礼，
@akanimax :)

BMSG-GAN 快速上手指南

BMSG-GAN 是 MSG-GAN（多尺度梯度生成对抗网络）的 PyTorch 实现，旨在解决传统 GAN 训练不稳定的问题。它通过允许判别器到生成器的多尺度梯度流动，实现同步的多尺度图像生成。

注意：本仓库仅为 PyTorch 复现版本。如需官方 TensorFlow 代码及预训练模型，请访问 msg-stylegan-tf。

1. 环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• Python: 3.6 或更高版本
• 深度学习框架: PyTorch (建议最新稳定版)
• 硬件: 推荐使用 NVIDIA GPU (实验中使用了 Tesla V100)，支持 CUDA

前置依赖安装：

建议使用 pip 安装必要的 Python 库。国内用户可使用清华源加速下载：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torch torchvision tqdm pillow numpy

2. 安装步骤

本项目无需复杂的编译安装，只需克隆仓库即可使用。

# 克隆仓库
git clone https://github.com/akanimax/BMSG-GAN.git

# 进入项目目录
cd BMSG-GAN

# 确认源代码位于 sourcecode/ 目录下
ls sourcecode/

3. 基本使用

核心参数建议

为了获得最佳效果，请遵循以下官方推荐的训练配置：

• 学习率: 生成器 (G) 和判别器 (D) 均设置为 0.003。
• 损失函数: 使用 relativistic-hinge (默认值)。

启动训练

进入 sourcecode/ 目录并运行 train.py 脚本。以下是一个训练 256x256 分辨率图像的示例命令：

cd sourcecode

python train.py --depth=7 \
                --latent_size=512 \
                --images_dir=<path_to_your_images> \
                --sample_dir=samples/exp_1 \
                --model_dir=models/exp_1 \
                --g_lr=0.003 \
                --d_lr=0.003 \
                --loss_function=relativistic-hinge \
                --batch_size=32 \
                --feedback_factor=500 \
                --checkpoint_factor=10

参数说明：

• --images_dir: 您的训练数据集图片路径。
• --depth: GAN 的深度，对应生成图像的分辨率（例如 depth=7 对应 256x256）。
• --sample_dir: 生成样本保存路径。
• --model_dir: 模型权重保存路径。
• --batch_size: 批大小，请根据显存大小调整。
• --flip_augment: 如需在训练中随机镜像图像，可添加此标志。

训练开始后，您将观察到从低分辨率到高分辨率的图像逐渐清晰并同步的过程。生成的样本将保存在指定的 sample_dir 中。