基于 DiffAugment 的数据高效 GAN

项目 | 论文 | 数据集 | 视频 | 演示文稿

仅使用 100 张奥巴马照片、生气的猫、熊猫、叹息桥、美第奇喷泉、天坛的照片生成，无需预训练。

[新增！] 现已提供基于 DiffAugment-stylegan2-pytorch 的 PyTorch 训练代码！

[新增！] 我们的 Colab 教程 已发布！

[新增！] 现已支持 FFHQ 数据集的训练！请参阅 DiffAugment-stylegan2 的 README。

[新增！] 使用 generate_gif.py 即可生成 100 张样本的插值视频！

[新增！] 我们的 DiffAugment-biggan-imagenet 仓库（用于 TPU 训练）已发布！

[新增！] 我们的 DiffAugment-biggan-cifar PyTorch 仓库已发布！

本仓库包含我们在 PyTorch 和 TensorFlow 中对可微增强（DiffAugment）的实现。它可用于显著提升 GAN 训练的数据效率。我们提供了 DiffAugment-stylegan2（TensorFlow）和 DiffAugment-stylegan2-pytorch（PyTorch），用于 GPU 训练；以及 DiffAugment-biggan-imagenet（TensorFlow），用于 TPU 训练。

无预训练条件下的低样本量生成。借助 DiffAugment，我们的模型仅需 100 张奥巴马肖像、生气的猫或熊猫等收集到的 100 张样本数据集，即可在 256×256 分辨率下生成高保真图像；同样地，仅用 AnimalFace 数据集中的 160 只猫或 389 只狗也能实现高质量生成。

CIFAR-10 上的无条件生成结果。StyleGAN2 在训练数据较少时性能会大幅下降。而借助 DiffAugment，我们仅使用 20% 的训练数据，便能大致匹配其 FID，并超越其 Inception Score（IS）。

用于数据高效 GAN 训练的可微增强
赵圣宇、刘志坚、林继、朱俊彦 和 韩松
麻省理工学院、清华大学、Adobe Research、卡内基梅隆大学
arXiv

概述

DiffAugment 用于更新判别器 D（左）和生成器 G（右）的概述。DiffAugment 将增强操作 T 同时应用于真实样本 x 和生成输出 G(z)。当更新 G 时，梯度需要通过 T 进行反向传播（iii），这要求 T 对输入是可微的。

使用 100 张图片进行训练与生成

要使用我们预训练的模型生成插值视频：

cd DiffAugment-stylegan2
python generate_gif.py -r mit-han-lab:DiffAugment-stylegan2-100-shot-obama.pkl -o obama.gif

或者训练一个新模型：

python run_low_shot.py --dataset=100-shot-obama --num-gpus=4

你也可以尝试 100-shot-grumpy_cat、100-shot-panda、100-shot-bridge_of_sighs、100-shot-medici_fountain、100-shot-temple_of_heaven、100-shot-wuzhen，或是放入你自己的训练图片文件夹。有关依赖项和详细信息，请参阅 DiffAugment-stylegan2 的 README。

[新增！] 现已提供 PyTorch 训练：

cd DiffAugment-stylegan2-pytorch
python train.py --outdir=training-runs --data=https://data-efficient-gans.mit.edu/datasets/100-shot-obama.zip --gpus=1

DiffAugment 用于 StyleGAN2

要在 100 张样本数据集、CIFAR、FFHQ 或 LSUN 上运行 StyleGAN2 + DiffAugment 进行无条件生成，请参阅 DiffAugment-stylegan2 的 README，或对于 PyTorch 版本，请参阅 DiffAugment-stylegan2-pytorch。

DiffAugment 用于 BigGAN

请参阅 DiffAugment-biggan-cifar 的 README，以运行 BigGAN + DiffAugment 进行 CIFAR 数据集上的条件生成（使用 GPU）；并参阅 DiffAugment-biggan-imagenet 的 README，以在 ImageNet 数据集上运行（使用 TPU）。

在你自己的训练中使用 DiffAugment

为了帮助你在自己的代码库中使用 DiffAugment，我们在 DiffAugment_tf.py 和 DiffAugment_pytorch.py 中提供了可移植的 TensorFlow 和 PyTorch 版本的 DiffAugment 操作。通常，只需将每个 D(x) 替换为 D(T(x))，即可轻松将 DiffAugment 应用于任何模型，其中 x 可以是真实图像或假图像，D 是判别器，T 是 DiffAugment 操作。例如，

from DiffAugment_pytorch import DiffAugment
# from DiffAugment_tf import DiffAugment
policy = 'color,translation,cutout' # 如果你的数据集像我们的那样小（例如，
# 只有几百张图片），我们建议使用最强的 Color + Translation + Cutout。
# 对于大型数据集，可以尝试使用 ['color', 'translation', 'cutout'] 中的部分变换。

# 欢迎探索更多 DiffAugment 变换！

...
# 训练循环：更新判别器 D
real_images = sample_real_images() # 一批真实图像
z = sample_latent_vectors()
fake_images = Generator(z) # 一批生成图像
real_scores = Discriminator(DiffAugment(real_images, policy=policy))
fake_scores = Discriminator(DiffAugment(fake_images, policy=policy))
# 根据 real_scores 和 fake_scores 计算判别器的损失...
...

...
# 训练循环：更新生成器 G
z = sample_latent_vectors()
fake_images = Generator(z) # 一批生成图像
fake_scores = Discriminator(DiffAugment(fake_images, policy=policy))
# 根据 fake_scores 计算生成器的损失...
...
我们实现了 Color、Translation 和 Cutout 这三种 DiffAugment 变换，效果如下所示：

<img src="https://oss.gittoolsai.com/images/mit-han-lab_data-efficient-gans_readme_e36a48e68a70.jpg" width="800px"/>

## 引用

如果您觉得这段代码有所帮助，请引用我们的论文：

@inproceedings{zhao2020diffaugment, title={Differentiable Augmentation for Data-Efficient GAN Training}, author={Zhao, Shengyu and Liu, Zhijian and Lin, Ji and Zhu, Jun-Yan and Han, Song}, booktitle={Conference on Neural Information Processing Systems (NeurIPS)}, year={2020} }

## 致谢

我们感谢 NSF 职业生涯奖 #1943349、MIT-IBM 沃森人工智能实验室、谷歌、Adobe 和索尼对本研究的支持。本研究还得到了谷歌 TensorFlow Research Cloud (TFRC) 提供的 Cloud TPU 的支持。同时，我们也感谢 William S. Peebles 和 Yijun Li 提出的有益建议。

Data-Efficient GANs (DiffAugment) 快速上手指南

本指南帮助开发者快速使用 Data-Efficient GANs，通过可微分数据增强（DiffAugment）技术，在极少样本（如仅 100 张图片）下训练高质量的生成对抗网络（GAN）。

环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu 18.04+) 或 macOS
• GPU: NVIDIA GPU (支持 CUDA)，显存建议 8GB 以上（根据模型大小调整）
• Python: 3.6 - 3.8

前置依赖

本项目提供 PyTorch 和 TensorFlow 两个版本，请根据需求选择安装基础框架：

• PyTorch 版本: 需要 torch, torchvision, ninja
• TensorFlow 版本: 需要 tensorflow-gpu (适用于 StyleGAN2 原版), numpy, scipy

提示: 若需训练 BigGAN-ImageNet 并使用 TPU，请参考官方仓库中的 DiffAugment-biggan-imagenet 分支说明。

安装步骤

克隆项目仓库并进入目录：

git clone https://github.com/mit-han-lab/data-efficient-gans.git
cd data-efficient-gans

方案 A：使用 PyTorch (推荐新手)

进入 PyTorch 子目录并安装依赖（需确保已安装 PyTorch）：

cd DiffAugment-stylegan2-pytorch
pip install -r requirements.txt

方案 B：使用 TensorFlow

进入 TensorFlow 子目录并安装依赖（需确保已安装 TensorFlow）：

cd DiffAugment-stylegan2
pip install -r requirements.txt

基本使用

以下示例展示如何使用 PyTorch 版本 进行训练和生成。

1. 快速训练 (100 张样本)

使用官方提供的 100 张奥巴马照片数据集进行训练（单卡 GPU）：

cd DiffAugment-stylegan2-pytorch
python train.py --outdir=training-runs --data=https://data-efficient-gans.mit.edu/datasets/100-shot-obama.zip --gpus=1

注：你可以将 --data 参数替换为其他内置数据集（如 100-shot-grumpy_cat.zip）或本地图片文件夹路径。

2. 使用预训练模型生成插值视频

下载预训练模型后，运行以下命令生成 GIF 动图：

cd DiffAugment-stylegan2
python generate_gif.py -r mit-han-lab:DiffAugment-stylegan2-100-shot-obama.pkl -o obama.gif

3. 集成到自己的代码中

若你想在自己的 GAN 项目中应用 DiffAugment，只需引入增强模块并在判别器输入前应用变换。

PyTorch 示例代码：

from DiffAugment_pytorch import DiffAugment

# 定义增强策略：小数据集推荐使用 'color,translation,cutout'
policy = 'color,translation,cutout'

# ... 在训练循环中 ...

# 更新判别器 (D)
reals = sample_real_images()
z = sample_latent_vectors()
fakes = Generator(z)

# 对真实图片和生成图片同时应用 DiffAugment
real_scores = Discriminator(DiffAugment(reals, policy=policy))
fake_scores = Discriminator(DiffAugment(fakes, policy=policy))

# 计算 D 的损失...

# 更新生成器 (G)
z = sample_latent_vectors()
fakes = Generator(z)

# 仅对生成图片应用 DiffAugment (梯度需回传通过增强操作)
fake_scores = Discriminator(DiffAugment(fakes, policy=policy))

# 计算 G 的损失...

TensorFlow 示例代码： 只需将导入语句改为 from DiffAugment_tf import DiffAugment，其余逻辑相同。