stylegan2-ada-pytorch

StyleGAN2-ADA — 官方 PyTorch 实现

使用有限数据训练生成对抗网络 Tero Karras, Miika Aittala, Janne Hellsten, Samuli Laine, Jaakko Lehtinen, Timo Aila https://arxiv.org/abs/2006.06676

摘要：使用过少数据训练生成对抗网络（GAN）通常会导致判别器过拟合，使训练发散。我们提出了一种自适应判别器增强机制，显著稳定了有限数据条件下的训练。这种方法无需更改损失函数或网络架构，既适用于从头开始训练，也适用于在另一个数据集上微调现有GAN。我们在多个数据集上证明，现在只需几千张训练图像就能获得良好结果，通常可以匹配使用多一个数量级图像的StyleGAN2结果。我们预计这将为GAN开辟新的应用领域。我们还发现，广泛使用的CIFAR-10实际上是一个有限数据基准，并将FID记录从5.59提升到2.42。

如需商业咨询，请访问我们的网站并提交表单：NVIDIA研究许可

发布说明

本仓库是在PyTorch中对StyleGAN2-ADA的忠实重新实现，专注于正确性、性能和兼容性。

正确性

• 全面支持所有主要训练配置。
• 对图像质量、训练曲线和质量指标与TensorFlow版本进行了广泛验证。
• 除伪随机数和浮点运算的影响外，预期结果在所有情况下都能匹配。

性能

• 在NVIDIA Tesla V100 GPU上，训练速度通常比TensorFlow版本快5%至30%。
• 高分辨率下的推理速度最高提升35%，但在低分辨率下可能略慢。
• GPU内存使用与TensorFlow版本相当。
• 训练新网络时启动时间更快（<50秒），使用预训练网络时也更快（<4秒）。
• 新增命令行选项用于调整训练性能。

兼容性

• 兼容使用TensorFlow版本创建的旧网络pickle文件。
• 新的基于ZIP/PNG的数据集格式，实现与现有第三方工具的最大互操作性。
• 不再支持TFRecords数据集 — 需要转换为新格式。
• 日志、指标和训练曲线采用新的基于JSON的格式。
• 如果安装了TensorBoard，训练曲线也会以旧的TFEvents格式导出。
• 命令行语法大部分保持不变，有少数例外（如dataset_tool.py）。
• 不支持比较方法（--cmethod、--dcap、--cfg=cifarbaseline、--aug=adarv）
• 截断现在默认禁用。

数据仓库

要求

• 支持Linux和Windows，但出于性能和兼容性考虑，我们推荐使用Linux。
• 1-8个高端NVIDIA GPU，至少12 GB内存。我们使用配备8个Tesla V100 GPU的NVIDIA DGX-1进行所有测试和开发。
• 64位Python 3.7和PyTorch 1.7.1。PyTorch安装说明请参见https://pytorch.org/。
• CUDA工具包11.0或更高版本。如果使用RTX 3090，请至少使用11.1版本。（为什么需要单独安装CUDA工具包？请参见#2中的评论。）
• Python库：pip install click requests tqdm pyspng ninja imageio-ffmpeg==0.4.3。我们使用Anaconda3 2020.11发行版，默认安装了大部分这些库。
• Docker用户：使用提供的Dockerfile构建包含所需库依赖的镜像。

代码大量依赖于使用NVCC即时编译的自定义PyTorch扩展。在Windows上，编译需要Microsoft Visual Studio。我们建议安装Visual Studio社区版，并使用"C:\Program Files (x86)\Microsoft Visual Studio\<VERSION>\Community\VC\Auxiliary\Build\vcvars64.bat"将其添加到PATH中。

入门

预训练网络存储为*.pkl文件，可以使用本地文件名或URL引用：

# 生成精选的MetFaces图像，无截断（图10左）
python generate.py --outdir=out --trunc=1 --seeds=85,265,297,849 \
    --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metfaces.pkl

# 生成未经筛选的MetFaces图像，带截断（图12左上）
python generate.py --outdir=out --trunc=0.7 --seeds=600-605 \
    --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metfaces.pkl

# 生成类条件CIFAR-10图像（图17左，汽车）
python generate.py --outdir=out --seeds=0-35 --class=1 \
    --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/cifar10.pkl

# 风格混合示例
python style_mixing.py --outdir=out --rows=85,100,75,458,1500 --cols=55,821,1789,293 \
    --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metfaces.pkl

上述命令的输出被放置在 out/*.png 下，由 --outdir 控制。下载的网络 pickle 文件被缓存在 $HOME/.cache/dnnlib 下，可以通过设置 DNNLIB_CACHE_DIR 环境变量来覆盖。默认的 PyTorch 扩展构建目录是 $HOME/.cache/torch_extensions，可以通过设置 TORCH_EXTENSIONS_DIR 来覆盖。

Docker：你可以使用 Docker 运行上述精选图像示例，如下所示：

docker build --tag sg2ada:latest .
./docker_run.sh python3 generate.py --outdir=out --trunc=1 --seeds=85,265,297,849 \
    --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metfaces.pkl

注意：Docker 镜像需要 NVIDIA 驱动版本 r455.23 或更高。

旧版网络：上述命令可以加载使用之前 TensorFlow 版本的 StyleGAN2 和 StyleGAN2-ADA 创建的大多数网络 pickle 文件。但是，为了未来的兼容性，我们建议将这些旧版 pickle 文件转换为 PyTorch 版本使用的新格式：

python legacy.py \
    --source=https://nvlabs-fi-cdn.nvidia.com/stylegan2/networks/stylegan2-cat-config-f.pkl \
    --dest=stylegan2-cat-config-f.pkl

将图像投影到潜在空间

要为给定的图像文件找到匹配的潜在向量，请运行：

python projector.py --outdir=out --target=~/mytargetimg.png \
    --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/ffhq.pkl

为获得最佳结果，目标图像应该被裁剪和对齐，类似于 FFHQ 数据集。上述命令保存投影目标 out/target.png、结果 out/proj.png、潜在向量 out/projected_w.npz 和进程视频 out/proj.mp4。你可以通过为 generate.py 指定 --projected_w 来渲染结果潜在向量：

python generate.py --outdir=out --projected_w=out/projected_w.npz \
    --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/ffhq.pkl

在 Python 中使用网络

你可以在自己的 Python 代码中使用预训练网络，如下所示：

with open('ffhq.pkl', 'rb') as f:
    G = pickle.load(f)['G_ema'].cuda()  # torch.nn.Module
z = torch.randn([1, G.z_dim]).cuda()    # 潜在编码
c = None                                # 类别标签（本例中未使用）
img = G(z, c)                           # NCHW，float32，动态范围 [-1, +1]

上述代码要求 torch_utils 和 dnnlib 可以通过 PYTHONPATH 访问。它不需要网络本身的源代码 —— 它们的类定义通过 torch_utils.persistence 从 pickle 文件中加载。

pickle 文件包含三个网络。'G' 和 'D' 是训练过程中的瞬时快照，'G_ema' 表示生成器权重在几个训练步骤中的移动平均值。这些网络是 torch.nn.Module 的常规实例，在导入时所有参数和缓冲区都放置在 CPU 上，默认情况下禁用梯度计算。

生成器由两个子模块组成，G.mapping 和 G.synthesis，可以单独执行。它们还支持各种附加选项：

w = G.mapping(z, c, truncation_psi=0.5, truncation_cutoff=8)
img = G.synthesis(w, noise_mode='const', force_fp32=True)

请参考 generate.py、style_mixing.py 和 projector.py 获取更多示例。

准备数据集

数据集以未压缩的 ZIP 归档形式存储，其中包含未压缩的 PNG 文件和用于标签的元数据文件 dataset.json。

可以从包含图像的文件夹创建自定义数据集；有关更多信息，请参阅 python dataset_tool.py --help。或者，也可以直接使用该文件夹作为数据集，无需先通过 dataset_tool.py 处理，但这样可能会导致性能不佳。

不支持旧版 TFRecords 数据集 —— 请参阅下面的说明了解如何转换它们。

FFHQ：

第 1 步：以 TFRecords 格式下载 Flickr-Faces-HQ 数据集。

第 2 步：使用 StyleGAN2-ADA 的 TensorFlow 版本 中的 dataset_tool.py 从 TFRecords 中提取图像：

# 使用 TensorFlow 版本的 dataset_tool.py
# https://github.com/NVlabs/stylegan2-ada/
python ../stylegan2-ada/dataset_tool.py unpack \
    --tfrecord_dir=~/ffhq-dataset/tfrecords/ffhq --output_dir=/tmp/ffhq-unpacked

第 3 步：使用本仓库中的 dataset_tool.py 创建 ZIP 归档：

# 原始 1024x1024 分辨率
python dataset_tool.py --source=/tmp/ffhq-unpacked --dest=~/datasets/ffhq.zip

# 缩小到 256x256 分辨率
#
# 注意：--resize-filter=box 是为了复现论文中显示的 FID 分数。
# 如果你不需要完全匹配，最好不要使用这个选项，默认使用 Lanczos。
# 参见 https://github.com/NVlabs/stylegan2-ada-pytorch/issues/283#issuecomment-1731217782
python dataset_tool.py --source=/tmp/ffhq-unpacked --dest=~/datasets/ffhq256x256.zip \
    --width=256 --height=256 --resize-filter=box

MetFaces：下载 MetFaces 数据集 并创建 ZIP 归档：

python dataset_tool.py --source=~/downloads/metfaces/images --dest=~/datasets/metfaces.zip

AFHQ：下载 AFHQ 数据集 并创建 ZIP 归档：

python dataset_tool.py --source=~/downloads/afhq/train/cat --dest=~/datasets/afhqcat.zip
python dataset_tool.py --source=~/downloads/afhq/train/dog --dest=~/datasets/afhqdog.zip
python dataset_tool.py --source=~/downloads/afhq/train/wild --dest=~/datasets/afhqwild.zip

CIFAR-10：下载 CIFAR-10 Python 版本 并转换为 ZIP 归档：

python dataset_tool.py --source=~/downloads/cifar-10-python.tar.gz --dest=~/datasets/cifar10.zip

LSUN：从 LSUN 项目页面 下载所需类别并转换为 ZIP 归档：

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

python dataset_tool.py --source=~/downloads/lsun/raw/car_lmdb --dest=~/datasets/lsuncar200k.zip \
    --transform=center-crop-wide --width=512 --height=384 --max_images=200000

BreCaHAD：

第 1 步：下载 BreCaHAD 数据集。

第 2 步：使用 StyleGAN2-ADA 的 TensorFlow 版本 中的 dataset_tool.py 提取 512x512 分辨率的裁剪：

# 使用 TensorFlow 版本的 dataset_tool.py
# https://github.com/NVlabs/stylegan2-ada/
python dataset_tool.py extract_brecahad_crops --cropsize=512 \
    --output_dir=/tmp/brecahad-crops --brecahad_dir=~/downloads/brecahad/images

第 3 步：使用本仓库中的 dataset_tool.py 创建 ZIP 归档：

python dataset_tool.py --source=/tmp/brecahad-crops --dest=~/datasets/brecahad.zip

训练新网络

最基本的形式，训练新网络可以归结为：

python train.py --outdir=~/training-runs --data=~/mydataset.zip --gpus=1 --dry-run
python train.py --outdir=~/training-runs --data=~/mydataset.zip --gpus=1

第一个命令是可选的；它验证参数，打印出训练配置，然后退出。第二个命令启动实际训练。

在这个例子中，结果保存在新创建的目录 ~/training-runs/<ID>-mydataset-auto1 中，由 --outdir 控制。训练过程会定期导出网络 pickle 文件（network-snapshot-<INT>.pkl）和示例图像（fakes<INT>.png）（由 --snap 控制）。对于每个 pickle 文件，它还会评估 FID（由 --metrics 控制）并将结果分数记录在 metric-fid50k_full.jsonl 中（如果安装了 TensorBoard，也会记录 TFEvents）。 输出目录的名称反映了训练配置。例如，00000-mydataset-auto1表示基本配置为auto1，意味着超参数是自动选择的，用于在一个GPU上训练。基本配置由--cfg控制：

训练配置可以通过其他命令行选项进一步自定义：

• --aug=noaug 禁用ADA。
• --cond=1 启用类条件训练（需要带标签的数据集）。
• --mirror=1 通过x轴翻转扩增数据集。通常有益，即使使用ADA也是如此。
• --resume=ffhq1024 --snap=10 从在1024x1024分辨率下训练的FFHQ进行迁移学习。
• --resume=~/training-runs/<NAME>/network-snapshot-<INT>.pkl 恢复之前的训练运行。
• --gamma=10 覆盖R1 gamma值。我们建议对每个新数据集尝试几个不同的值。
• --aug=ada --target=0.7 调整ADA目标值（默认：0.6）。
• --augpipe=blit 启用像素混合但禁用所有其他增强。
• --augpipe=bgcfnc 启用所有可用的增强（混合、几何、颜色、滤镜、噪声、裁剪）。

请参考python train.py --help获取完整列表。

预期训练时间

总训练时间严重依赖于分辨率、GPU数量、数据集、期望质量和超参数。下表列出了达到训练不同阶段的预期壁钟时间，以向判别器展示的真实图像数量（以千计，"kimg"）来衡量：

上述测量使用NVIDIA Tesla V100 GPU和默认设置（--cfg=auto --aug=ada --metrics=fid50k_full）完成。"秒/kimg"显示了原始训练性能的预期变化范围，如log.txt中报告的那样。"GPU内存"和"CPU内存"显示了观察到的最高内存消耗，不包括由torch.backends.cudnn.benchmark引起的开始时的峰值。

在典型情况下，需要25000 kimg或更多才能达到收敛，但在5000 kimg左右结果已经相当不错了。1000 kimg通常足以进行迁移学习，迁移学习往往收敛得更快。下图显示了不同数据集的示例收敛曲线，以壁钟时间为函数，使用与上述相同的设置：

注意：--cfg=auto作为超参数的合理首选猜测，但不一定能为给定数据集带来最佳结果。例如，对于1024x1024分辨率的FFHQ-140k，--cfg=stylegan2产生的FID明显优于上图所示。我们建议为每个新数据集至少尝试几个不同的--gamma值。

质量指标

默认情况下，train.py会自动计算训练过程中导出的每个网络pickle的FID。我们建议定期检查metric-fid50k_full.jsonl（或TensorBoard）以监控训练进度。如果需要，可以使用--metrics=none禁用自动计算，以略微加快训练速度（3%–9%）。

训练后还可以计算其他质量指标：

# 之前的训练运行：自动查找选项，将结果保存到JSONL文件。
python calc_metrics.py --metrics=pr50k3_full \
    --network=~/training-runs/00000-ffhq10k-res64-auto1/network-snapshot-000000.pkl

# 预训练网络pickle：明确指定数据集，将结果打印到stdout。
python calc_metrics.py --metrics=fid50k_full --data=~/datasets/ffhq.zip --mirror=1 \
    --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/ffhq.pkl

第一个示例查找训练配置并执行与在训练期间指定--metrics=pr50k3_full相同的操作。第二个示例下载预训练的网络pickle，在这种情况下必须明确指定--mirror和--data的值。

请注意，许多指标在首次为新数据集计算时都有显著的一次性成本（最多30分钟）。另外请注意，评估每次都使用不同的随机种子，因此如果多次计算相同的指标，结果会有所不同。

我们在ADA论文中采用以下指标。执行时间和GPU内存使用情况是针对一个NVIDIA Tesla V100 GPU在1024x1024分辨率下报告的：

此外，还支持以下来自StyleGAN和StyleGAN2论文的指标：

参考文献：

1. GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium，Heusel等人，2017
2. Demystifying MMD GANs，Bińkowski等人，2018
3. Improved Precision and Recall Metric for Assessing Generative Models，Kynkäänniemi等人，2019
4. Improved Techniques for Training GANs，Salimans等人，2016
5. A Style-Based Generator Architecture for Generative Adversarial Networks，Karras等人，2018

许可证

版权所有 © 2021，NVIDIA Corporation。保留所有权利。

本作品根据Nvidia Source Code License提供。

引用

@inproceedings{Karras2020ada,
  title     = {Training Generative Adversarial Networks with Limited Data},
  author    = {Tero Karras and Miika Aittala and Janne Hellsten and Samuli Laine and Jaakko Lehtinen and Timo Aila},
  booktitle = {Proc. NeurIPS},
  year      = {2020}
}

开发

这是一个研究参考实现，被视为一次性代码发布。因此，我们不接受以拉取请求形式提交的外部代码贡献。

致谢

我们感谢David Luebke提供的有益评论；感谢Tero Kuosmanen和Sabu Nadarajan在计算基础设施方面的支持；感谢Edgar Schönfeld在设置无条件BigGAN方面的指导。