Diffusion-GAN —— 官方 PyTorch 实现

Diffusion-GAN：使用扩散过程训练生成对抗网络（GANs）
Zhendong Wang, Huangjie Zheng, Pengcheng He, Weizhu Chen 和 Mingyuan Zhou
https://arxiv.org/abs/2206.02262

摘要：为了稳定训练生成对抗网络（Generative Adversarial Networks, GANs），向判别器输入中注入实例噪声（instance noise）被认为是一种理论上合理的解决方案，然而在实践中尚未充分兑现其潜力。本文提出 Diffusion-GAN，它采用一个定义在前向扩散链（forward diffusion chain）所有扩散步骤上的高斯混合分布（Gaussian mixture distribution）来注入实例噪声。从该混合分布中随机采样的样本（由真实或生成数据经扩散得到）将作为判别器的输入。生成器通过前向扩散链反向传播梯度进行更新，而扩散链长度会自适应调整，以控制每个训练步骤允许的最大噪声-数据比。理论分析验证了所提出的 Diffusion-GAN 的合理性，它提供了一种与模型和领域无关的可微增强方法。在多个数据集上丰富的实验表明，Diffusion-GAN 能够实现稳定且数据高效的 GAN 训练，在合成逼真图像任务上持续优于强大的 GAN 基线模型。

待办事项

• 初始代码发布
• 提供预训练模型

构建你的 Diffusion-GAN

在此，我们解释如何使用扩散过程训练通用 GAN。我们提供两种方式： a. 作为数据增强方法简单插件式使用； b. 在扩散链上训练 GAN，并使用依赖时间步（timestep-dependent）的判别器。 目前，我们尚未发现这两种方法在经验上有显著差异，但第二种方法具有更强的理论保证。我们推测，当判别器中应用更先进的时间步相关结构时，第二种方法可能表现更优，这留待未来研究。

简单插件式用法

• 根据 diffusion.py 文件设计合适的扩散过程（示例）
• 对判别器输入应用扩散：

logits = Discriminator(Diffusion(gen/real_images))

• 在训练迭代中加入扩散过程的自适应机制：

if update_diffusion:  # batch_idx % ada_interval == 0
    adjust = np.sign(sign(Discriminator(real_images)) - ada_target) * C  # C = (batch_size * ada_interval) / (ada_kimg * 1000)
    diffusion.p = (diffusion.p + adjust).clip(min=0., max=1.)
    diffusion.update_T()

完整版本

• 将扩散时间步 t 作为判别器的输入：logits = Discriminator(images, t)。 你可能需要对判别器架构进行一些修改。
• 其余步骤与“简单插件式”相同。注意由于判别器依赖于时间步，你需要收集 t。

diffused_images, t = Diffusion(images)
logits = Discrimnator(diffused_images, t)

训练我们的 Diffusion-GAN

环境要求

每个文件夹中我们都提供了 environment.yml 用于配置运行环境。注意，为适配最新版 PyTorch，我对 torch_utils 文件夹做了一些修改。由于 PyTorch 版本不同，性能可能会有轻微波动。所提供的检查点是在继承自 StyleGAN2-ADA 的旧环境中训练的，旧环境详情如下：

• 64 位 Python 3.7 和 PyTorch 1.7.1/1.8.1。PyTorch 安装说明请参见 https://pytorch.org/。
• CUDA 工具包 11.0 或更高版本。
• Python 库：pip install click requests tqdm pyspng ninja imageio-ffmpeg==0.4.3。

数据准备

在论文中，我们在以下数据集上训练了模型：CIFAR-10 (32 x 32)、STL-10 (64 x 64)、 LSUN (256 x 256)、AFHQ (512 x 512) 和 FFHQ (1024 x 1024)。 你可以从各自网站下载我们在论文中使用的数据集。要准备指定分辨率的数据集，例如运行：

python dataset_tool.py --source=~/downloads/lsun/raw/bedroom_lmdb --dest=~/datasets/lsun_bedroom200k.zip \
    --transform=center-crop --width=256 --height=256 --max_images=200000

python dataset_tool.py --source=~/downloads/lsun/raw/church_lmdb --dest=~/datasets/lsun_church200k.zip \
    --transform=center-crop-wide --width=256 --height=256 --max_images=200000

训练

我们在下方展示了所使用的训练命令。在大多数情况下，训练命令类似，因此以下以 CIFAR-10 数据集为例：

对于 Diffusion-GAN：

python train.py --outdir=training-runs --data="~/cifar10.zip" --gpus=4 --cfg cifar --kimg 50000 --aug no --target 0.6 --noise_sd 0.05 --ts_dist priority

对于 Diffusion-ProjectedGAN：

python train.py --outdir=training-runs --data="~/cifar10.zip" --gpus=4 --batch 64 --batch-gpu=16 --cfg fastgan --kimg 50000 --target 0.45 --d_pos first --noise_sd 0.5

对于 Diffusion-InsGen：

python train.py --outdir=training-runs --data="~/afhq-wild.zip" --gpus=8 --cfg paper512 --kimg 25000

我们遵循 StyleGAN2-ADA 的 config 设置，更多细节请参考其项目。其他主要超参数如下并作简要说明：

• --target：判别器目标值，用于平衡扩散强度。
• --aug：领域特定的图像增强（如 ADA 和可微分增强），用于评估与扩散的互补性。
• --noise_sd：扩散噪声标准差，在我们的实验中设为 0.05。
• --ts_dist：时间步采样分布 $\pi(t)$（论文中的符号）。

我们评估了两种 t 采样分布：['priority', 'uniform']， 其中 'priority' 表示论文中的公式 (11)，而 'uniform' 表示随机采样。在多数情况下，priority 略优；但在某些数据集（如 FFHQ）上，'uniform' 表现更好。

使用我们的检查点进行采样与评估

我们在 Hugging Face 页面提供了 Diffusion-GAN 的预训练检查点，下方为快速访问链接：

模型	数据集	分辨率	FID	检查点
Diffusion-StyleGAN2	CIFAR-10	32x32	3.19	下载
Diffusion-StyleGAN2	CelebA	64x64	1.69	下载
Diffusion-StyleGAN2	STL-10	64x64	11.53	下载
Diffusion-StyleGAN2	LSUN-Bedroom	256x256	3.65	下载
Diffusion-StyleGAN2	LSUN-Church	256x256	3.17	下载
Diffusion-StyleGAN2	FFHQ	1024x1024	2.83	下载
Diffusion-ProjectedGAN	CIFAR-10	32x32	2.54	下载
Diffusion-ProjectedGAN	STL-10	64x64	6.91	下载
Diffusion-ProjectedGAN	LSUN-Bedroom	256x256	1.43	下载
Diffusion-ProjectedGAN	LSUN-Church	256x256	1.85	下载
Diffusion-InsGen	AFHQ-Cat	512x512	2.40	下载
Diffusion-InsGen	AFHQ-Dog	512x512	4.83	下载
Diffusion-InsGen	AFHQ-Wild	512x512	1.51	下载

要生成样本，请运行以下命令：

# 使用预训练的 Diffusion-StyleGAN2 生成 FFHQ 样本
python generate.py --outdir=out --seeds=1-100 \
    --network=https://tsciencescu.blob.core.windows.net/projectshzheng/DiffusionGAN/diffusion-stylegan2-ffhq.pkl

# 使用预训练的 Diffusion-ProjectedGAN 生成 LSUN-Church 样本
python gen_images.py --outdir=out --seeds=1-100 \
    --network=https://tsciencescu.blob.core.windows.net/projectshzheng/DiffusionGAN/diffusion-projectedgan-lsun-church.pkl

上述命令中的检查点路径可替换为表格中任意预训练 Diffusion-GAN 检查点。

同样，可通过以下命令计算指标：

# 预训练网络模型：显式指定数据集，结果输出至标准输出
python calc_metrics.py --metrics=fid50k_full --data=~/datasets/ffhq.zip --mirror=1 \
    --network=https://tsciencescu.blob.core.windows.net/projectshzheng/DiffusionGAN/diffusion-stylegan2-ffhq.pkl

引用

@article{wang2022diffusiongan,
  title     = {Diffusion-GAN: Training GANs with Diffusion},
  author    = {Wang, Zhendong and Zheng, Huangjie and He, Pengcheng and Chen, Weizhu and Zhou, Mingyuan},
  journal   = {arXiv preprint arXiv:2206.02262},
  year      = {2022},
  url       = {https://arxiv.org/abs/2206.02262}
}

致谢

我们的代码基于 Karras 等人的优秀 StyleGAN2-ADA 仓库、Ceyuan Yang 等人的 InsGen 仓库 以及 Axel Sauer 等人的 ProjectedGAN 仓库 构建。

Diffusion-GAN 快速上手指南

环境准备

系统要求

• 64位 Python 3.7
• PyTorch 1.7.1 或 1.8.1（推荐使用官方源或清华/阿里云镜像加速）
• CUDA Toolkit 11.0 或更高版本（需与 PyTorch 版本兼容）
• NVIDIA GPU（建议至少 8GB 显存）

前置依赖

pip install click requests tqdm pyspng ninja imageio-ffmpeg==0.4.3

国内用户建议配置 pip 镜像源（如清华源）加速安装：

pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

1. 克隆项目仓库：

git clone https://github.com/Zhendong-Wang/Diffusion-GAN.git
cd Diffusion-GAN

2. 创建并激活 Conda 环境（推荐）：

conda env create -f environment.yml
conda activate diffusion-gan

3. 编译必要组件（如有）：

# 根据具体子目录需求，例如进入 diffusion-stylegan2 目录后执行
cd diffusion-stylegan2
python setup.py build_ext --inplace

基本使用

数据准备示例（以 LSUN 卧室数据集为例）：

python dataset_tool.py --source=~/downloads/lsun/raw/bedroom_lmdb --dest=~/datasets/lsun_bedroom200k.zip \
    --transform=center-crop --width=256 --height=256 --max_images=200000

训练命令示例（CIFAR-10）：

python train.py --outdir=training-runs --data="~/cifar10.zip" --gpus=4 --cfg cifar --kimg 50000 --aug no --target 0.6 --noise_sd 0.05 --ts_dist priority

使用预训练模型生成图像（FFHQ 示例）：

python generate.py --outdir=out --seeds=1-100 \
    --network=https://tsciencescu.blob.core.windows.net/projectshzheng/DiffusionGAN/diffusion-stylegan2-ffhq.pkl

预训练模型也可从 Hugging Face 下载（国内用户可尝试代理或镜像加速）： https://huggingface.co/zhendongw/diffusion-gan/tree/main/checkpoints