无混叠生成对抗网络（StyleGAN3）
_{NeurIPS 2021 论文的官方 PyTorch 实现}

无混叠生成对抗网络
特罗·卡拉斯、米卡·艾塔拉、萨穆利·莱内、埃里克·哈尔科宁、扬内·赫尔斯特恩、雅各·莱蒂宁、蒂莫·艾拉
https://nvlabs.github.io/stylegan3

摘要：我们观察到，尽管典型的生成对抗网络具有分层卷积结构，但其合成过程却以一种不健康的方式依赖于绝对像素坐标。这表现为细节似乎被“粘”在图像坐标上，而非所描绘物体的表面。我们将根本原因归结为粗心的信号处理，导致生成器网络中出现混叠现象。通过将网络中的所有信号视为连续信号，我们推导出通用且微小的架构改进，确保有害信息不会渗入分层合成过程。由此产生的网络在 FID 指标上与 StyleGAN2 相当，但在内部表示上却有显著差异，并且即使在亚像素尺度下也完全具备平移和旋转等变性。我们的研究为更适合视频和动画的生成模型开辟了道路。

如需商务合作，请访问我们的官网并提交表格：NVIDIA Research Licensing

发布说明

本仓库是 stylegan2-ada-pytorch 的更新版本，新增多项功能：

• 无混叠生成器架构及训练配置（stylegan3-t、stylegan3-r）。
• 交互式可视化工具（visualizer.py）、频谱分析工具（avg_spectra.py）和视频生成工具（gen_video.py）。
• 等变性指标（eqt50k_int、eqt50k_frac、eqr50k）。
• 其他改进：降低内存占用、略微加快训练速度、修复若干 bug。

兼容性：

• 兼容使用 stylegan2-ada 和 stylegan2-ada-pytorch 创建的旧版网络快照。（注意：在 StyleGAN3 代码上运行旧版 StyleGAN2 模型，结果将与在 stylegan2-ada/stylegan2-ada-pytorch 上运行时相同。若要享受 StyleGAN3 架构的优势，需重新训练。）
• 支持旧版 StyleGAN2 的训练配置，包括 ADA 和迁移学习。详情请参阅 训练配置。
• 对 Ampere 架构 GPU 以及较新版本的 PyTorch、CuDNN 等的兼容性有所提升。

合成图像检测

虽然新的生成器方法带来了全新的媒体合成能力，但也可能对用于检测和溯源合成媒体的人工智能取证算法构成新的挑战。我们与参与 DARPA SemaFor 项目的数字取证研究人员合作，整理了一个合成图像数据集，使研究人员能够在公开发布前测试并验证其图像检测器的性能。更多详情请参见 此处。

补充资料

• 结果视频
• 精选示例图片
• StyleGAN3 预训练模型，适用于 T 配置（平移等变）和 R 配置（平移和旋转等变）。

  _{可通过 https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/<MODEL> 访问单个模型，其中 <MODEL> 包括：
  stylegan3-t-ffhq-1024x1024.pkl、stylegan3-t-ffhqu-1024x1024.pkl、stylegan3-t-ffhqu-256x256.pkl
  stylegan3-r-ffhq-1024x1024.pkl、stylegan3-r-ffhqu-1024x1024.pkl、stylegan3-r-ffhqu-256x256.pkl
  stylegan3-t-metfaces-1024x1024.pkl、stylegan3-t-metfacesu-1024x1024.pkl
  stylegan3-r-metfaces-1024x1024.pkl、stylegan3-r-metfacesu-1024x1024.pkl
  stylegan3-t-afhqv2-512x512.pkl
  stylegan3-r-afhqv2-512x512.pkl}
  

• StyleGAN2 预训练模型，与本代码库兼容。

  _{可通过 https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/<MODEL> 访问单个模型，其中 <MODEL> 包括：
  stylegan2-ffhq-1024x1024.pkl、stylegan2-ffhq-512x512.pkl、stylegan2-ffhq-256x256.pkl
  stylegan2-ffhqu-1024x1024.pkl、stylegan2-ffhqu-256x256.pkl
  stylegan2-metfaces-1024x1024.pkl、stylegan2-metfacesu-1024x1024.pkl
  stylegan2-afhqv2-512x512.pkl
  stylegan2-afhqcat-512x512.pkl、stylegan2-afhqdog-512x512.pkl、stylegan2-afhqwild-512x512.pkl
  stylegan2-brecahad-512x512.pkl、stylegan2-cifar10-32x32.pkl
  stylegan2-celebahq-256x256.pkl、stylegan2-lsundog-256x256.pkl}
  

需求

• 支持 Linux 和 Windows，但出于性能和兼容性考虑，我们推荐使用 Linux。
• 1–8 张高端 NVIDIA GPU，每张显存至少 12 GB。我们所有的测试和开发都是在 Tesla V100 和 A100 GPU 上完成的。
• 64 位 Python 3.8 和 PyTorch 1.9.0（或更高版本）。PyTorch 的安装说明请参见 https://pytorch.org。
• CUDA 工具包 11.1 或更高版本。（为什么需要单独安装 CUDA 工具包？请参阅 故障排除。）
• GCC 7 或更高版本（Linux）或 Visual Studio（Windows）编译器。推荐的 GCC 版本取决于 CUDA 版本，例如，请参阅 CUDA 11.4 系统要求。
• Python 库：确切的库依赖关系请参见 environment.yml。您可以使用 Miniconda3 中的以下命令来创建并激活 StyleGAN3 的 Python 环境：
  • conda env create -f environment.yml
  • conda activate stylegan3
• Docker 用户：
  • 确保已正确安装 NVIDIA 容器运行时。
  • 使用 提供的 Dockerfile 来构建包含所需库依赖的镜像。

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

有关常见安装和运行时问题的帮助，请参阅 故障排除。

快速入门

预训练网络以 *.pkl 文件的形式存储，可以通过本地文件名或 URL 引用：

# 使用预训练的 AFHQv2 模型生成一张图像（图 1 左侧的“Ours”）。
python gen_images.py --outdir=out --trunc=1 --seeds=2 \
    --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl

# 渲染一个由种子 0 至 31 构成的 4×2 插值网格。
python gen_video.py --output=lerp.mp4 --trunc=1 --seeds=0-31 --grid=4x2 \
    --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl

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

Docker：您可以通过以下方式使用 Docker 运行上述示例：

# 构建 stylegan3:latest 镜像
docker build --tag stylegan3 .

# 使用 Docker 运行 gen_images.py 脚本：
docker run --gpus all -it --rm --user $(id -u):$(id -g) \
    -v `pwd`:/scratch --workdir /scratch -e HOME=/scratch \
    stylegan3 \
    python gen_images.py --outdir=out --trunc=1 --seeds=2 \
         --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl

注意：Docker 镜像需要 NVIDIA 驱动程序版本 r470 或更高版本。

docker run 命令可能看起来有些复杂，下面我们来逐项解释：

• --gpus all -it --rm --user $(id -u):$(id -g)：启用所有 GPU，以当前用户的 UID/GID 运行交互式会话，避免 Docker 以 root 用户身份写入文件。
• -v pwd:/scratch --workdir /scratch：将当前工作目录（例如，您主机上的这个 Git 仓库根目录）挂载到容器内的 /scratch 目录，并将其设为当前工作目录。
• -e HOME=/scratch：让 PyTorch 和 StyleGAN3 代码知道临时文件（如预训练模型和自定义 PyTorch 扩展的构建结果）应该缓存在哪里。注意：如果您需要更精细的控制，可以分别设置 TORCH_EXTENSIONS_DIR（用于自定义扩展的构建目录）和 DNNLIB_CACHE_DIR（用于预训练模型的下载缓存）。这些缓存目录应位于持久化卷上，以便在多次 docker run 调用之间保留内容。

交互式可视化

此版本包含一个交互式模型可视化工具，可用于探索训练好的模型的各种特性。要启动它，请运行：

python visualizer.py

在 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)

完整的代码示例请参阅 gen_images.py。

准备数据集

数据集以未压缩的 ZIP 存档形式存储，其中包含未压缩的 PNG 文件和一个用于标签的元数据文件 dataset.json。自定义数据集可以从包含图像的文件夹中创建；更多信息请参阅 python dataset_tool.py --help。或者，也可以直接将文件夹用作数据集，无需先通过 dataset_tool.py 处理，但这样做可能会导致性能不佳。

FFHQ：下载 Flickr-Faces-HQ 数据集，格式为 1024×1024 像素的图像，并使用 dataset_tool.py 创建一个 ZIP 存档：

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

# 缩小至256x256分辨率。
python dataset_tool.py --source=/tmp/images1024x1024 --dest=~/datasets/ffhq-256x256.zip \
    --resolution=256x256

有关如何获取未对齐的FFHQ数据集图像的信息，请参阅FFHQ README。使用上述相同步骤创建用于训练和验证的ZIP存档。

MetFaces：下载MetFaces数据集，并创建一个ZIP存档：

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

有关如何获取未对齐的MetFaces数据集图像的信息，请参阅MetFaces README。使用上述相同步骤创建用于训练和验证的ZIP存档。

AFHQv2：下载AFHQv2数据集，并创建一个ZIP存档：

python dataset_tool.py --source=~/downloads/afhqv2 --dest=~/datasets/afhqv2-512x512.zip

请注意，上述命令会使用所有三个类别（猫、狗和野生动物）的所有图像创建一个单一的合并数据集，这与StyleGAN3论文中使用的设置一致。或者，您也可以为每个类别单独创建一个数据集：

python dataset_tool.py --source=~/downloads/afhqv2/train/cat --dest=~/datasets/afhqv2cat-512x512.zip
python dataset_tool.py --source=~/downloads/afhqv2/train/dog --dest=~/datasets/afhqv2dog-512x512.zip
python dataset_tool.py --source=~/downloads/afhqv2/train/wild --dest=~/datasets/afhqv2wild-512x512.zip

训练

您可以使用train.py来训练新网络。例如：

# 使用8个GPU为AFHQv2训练StyleGAN3-T。
python train.py --outdir=~/training-runs --cfg=stylegan3-t --data=~/datasets/afhqv2-512x512.zip \
    --gpus=8 --batch=32 --gamma=8.2 --mirror=1

# 使用1个GPU为MetFaces-U微调StyleGAN3-R，从预训练的FFHQ-U pickle开始。
python train.py --outdir=~/training-runs --cfg=stylegan3-r --data=~/datasets/metfacesu-1024x1024.zip \
    --gpus=8 --batch=32 --gamma=6.6 --mirror=1 --kimg=5000 --snap=5 \
    --resume=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-ffhqu-1024x1024.pkl

# 使用8个GPU在1024x1024分辨率下为FFHQ训练StyleGAN2。
python train.py --outdir=~/training-runs --cfg=stylegan2 --data=~/datasets/ffhq-1024x1024.zip \
    --gpus=8 --batch=32 --gamma=10 --mirror=1 --aug=noaug

请注意，结果质量和训练时间在很大程度上取决于具体的选项设置。其中最重要的几个选项（--gpus、--batch和--gamma）必须明确指定，并且应谨慎选择。有关完整选项列表，请参阅python train.py --help，有关一般指导原则和建议以及不同场景下的预期训练速度和内存使用情况，请参阅训练配置。

每次训练的结果都会保存到一个新创建的目录中，例如~/training-runs/00000-stylegan3-t-afhqv2-512x512-gpus8-batch32-gamma8.2。训练循环会定期导出网络pickle文件（network-snapshot-<KIMG>.pkl）和随机图像网格（fakes<KIMG>.png），间隔由--snap控制。对于每个导出的pickle文件，它会评估FID（由--metrics控制），并将结果记录在metric-fid50k_full.jsonl中。此外，它还会将各种统计信息记录在training_stats.jsonl中，如果安装了TensorBoard，则还会生成*.tfevents文件。

质量指标

默认情况下，train.py会在训练过程中自动计算每个网络pickle文件的FID值。我们建议定期检查metric-fid50k_full.jsonl（或TensorBoard）以监控训练进度。如果需要，可以通过设置--metrics=none来禁用自动计算，从而略微加快训练速度。

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

# 上一次训练运行：自动查找选项，将结果保存到JSONL文件。
python calc_metrics.py --metrics=eqt50k_int,eqr50k \
    --network=~/training-runs/00000-stylegan3-r-mydataset/network-snapshot-000000.pkl

# 预训练的网络pickle：显式指定数据集，将结果打印到标准输出。
python calc_metrics.py --metrics=fid50k_full --data=~/datasets/ffhq-1024x1024.zip --mirror=1 \
    --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-t-ffhq-1024x1024.pkl

第一个示例会查找训练配置，并执行与在训练时指定--metrics=eqt50k_int,eqr50k相同的操作。第二个示例则下载了一个预训练的网络pickle，在这种情况下，必须显式指定--data和--mirror的值。

请注意，计算这些指标可能相当耗时（最长可达1小时），并且许多指标对于每个新数据集都需要额外的一次性成本（最长可达30分钟）。此外，每次评估都会使用不同的随机种子，因此如果多次计算同一指标，结果可能会有所不同。

推荐的指标：

• fid50k_full：针对整个数据集的Fréchet inception距离^[1]。
• kid50k_full：针对整个数据集的Kernel inception距离^[2]。
• pr50k3_full：针对整个数据集的精确率和召回率^[3]。
• ppl2_wend：在W空间中，端点与完整图像之间的感知路径长度^[4]。
• eqt50k_int：相对于整数平移的等变性^[5]（EQ-T）。
• eqt50k_frac：相对于分数平移的等变性（EQ-T_frac）。
• eqr50k：相对于旋转的等变性（EQ-R）。

旧版指标：

• fid50k：针对5万张真实图像的Fréchet inception距离。
• kid50k：针对5万张真实图像的Kernel inception距离。
• pr50k3：针对5万张真实图像的精确率和召回率。
• is50k：CIFAR-10的Inception评分^[6]。

参考文献：

1. 通过双时间尺度更新规则训练的GAN收敛到局部纳什均衡，Heusel等人，2017年
2. 揭秘MMD GAN，Bińkowski等人，2018年
3. 用于评估生成模型的改进型精确率和召回率指标，Kynkäänniemi等人，2019年
4. 基于风格的生成器架构用于生成对抗网络，Karras等人，2018年
5. 无混叠生成对抗网络，Karras等人，2021年
6. 训练GAN的改进技术，Salimans等人，2016年

光谱分析

检查给定生成器光谱特性的最简单方法是使用visualizer.py中的内置FFT模式。此外，您还可以按照以下方式可视化平均2D功率谱（附录A，图15）：

# 计算数据集的均值和标准差，后续步骤需要用到。
python avg_spectra.py stats --source=~/datasets/ffhq-1024x1024.zip

# 计算训练数据的平均频谱。
python avg_spectra.py calc --source=~/datasets/ffhq-1024x1024.zip \
    --dest=tmp/training-data.npz --mean=112.684 --std=69.509

# 计算预训练生成器的平均频谱。
python avg_spectra.py calc \
    --source=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-ffhq-1024x1024.pkl \
    --dest=tmp/stylegan3-r.npz --mean=112.684 --std=69.509 --num=70000

# 展示结果。
python avg_spectra.py heatmap tmp/training-data.npz
python avg_spectra.py heatmap tmp/stylegan3-r.npz
python avg_spectra.py slices tmp/training-data.npz tmp/stylegan3-r.npz

许可证

版权所有 © 2021，英伟达公司及其关联公司。保留所有权利。

本作品根据 NVIDIA 源代码许可证 提供。

引用

@inproceedings{Karras2021,
  author = {Tero Karras 和 Miika Aittala 和 Samuli Laine 和 Erik H\"ark\"onen 和 Janne Hellsten 和 Jaakko Lehtinen 和 Timo Aila},
  title = {无混叠生成对抗网络},
  booktitle = {NeurIPS 会议论文集},
  year = {2021}
}

开发

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

致谢

我们感谢 David Luebke、Ming-Yu Liu、Koki Nagano、Tuomas Kynkäänniemi 和 Timo Viitanen 对早期草稿的审阅及有益建议。感谢 Frédo Durand 的早期讨论。感谢 Tero Kuosmanen 维护我们的计算基础设施。感谢 AFHQ 团队提供其数据集的更新版本。感谢 Getty Images 提供 Beaches 数据集中用于训练的图像。本项目未获得任何外部资助或额外收入。

StyleGAN3 快速上手指南

StyleGAN3 是 NVIDIA 推出的无混叠生成对抗网络，解决了传统 GAN 中细节“粘”在像素坐标上的问题，实现了完美的平移和旋转等变性，特别适合视频生成和动画应用。

环境准备

系统要求：

• 操作系统：推荐 Linux（性能与兼容性最佳），支持 Windows。
• GPU：1-8 张高端 NVIDIA 显卡，显存至少 12GB（测试基于 Tesla V100/A100）。
• 软件版本：
  • Python 3.8 (64-bit)
  • PyTorch 1.9.0 或更高版本
  • CUDA Toolkit 11.1 或更高版本
  • 编译器：Linux 需 GCC 7+，Windows 需 Visual Studio（用于编译自定义 PyTorch 扩展）。

依赖安装： 建议使用 Miniconda3 创建独立环境。虽然官方未提供国内镜像源配置，但在使用 conda 时可手动替换为清华或中科大源以加速下载。

# 创建并激活环境
conda env create -f environment.yml
conda activate stylegan3

注意：代码依赖即时编译的自定义 PyTorch 扩展。Windows 用户需确保已安装 Visual Studio 并将其 vcvars64.bat 加入 PATH。

安装步骤

克隆仓库后无需额外编译命令，首次运行脚本时会自动编译所需的 CUDA 扩展。

git clone https://github.com/NVlabs/stylegan3.git
cd stylegan3

若使用 Docker（推荐，可避免环境配置问题）：

# 构建镜像
docker build --tag stylegan3 .

# 运行容器（挂载当前目录，启用 GPU）
docker run --gpus all -it --rm --user $(id -u):$(id -g) \
    -v $(pwd):/scratch --workdir /scratch -e HOME=/scratch \
    stylegan3

基本使用

1. 生成图像

使用预训练模型（如 AFHQv2 动物数据集）生成单张图片。模型文件会自动下载并缓存。

python gen_images.py --outdir=out --trunc=1 --seeds=2 \
    --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl

• 输出图片将保存在 out/ 目录下。
• stylegan3-r 表示支持平移和旋转等变性的模型配置。

2. 生成插值视频

生成种子点之间的平滑过渡视频。

python gen_video.py --output=lerp.mp4 --trunc=1 --seeds=0-31 --grid=4x2 \
    --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl

3. 交互式可视化

启动可视化工具，实时探索潜空间（Latent Space）特性。

python visualizer.py

4. 在 Python 代码中调用

直接在脚本中加载预训练模型进行推理：

import torch
import pickle

# 加载模型
with open('ffhq.pkl', 'rb') as f:
    G = pickle.load(f)['G_ema'].cuda()  # 使用指数移动平均生成器

# 准备潜变量
z = torch.randn([1, G.z_dim]).cuda()    
c = None                                # 类别标签（本例不使用）

# 生成图像
img = G(z, c)                           # 输出范围 [-1, +1]

提示：运行上述 Python 代码前，请确保当前目录包含 torch_utils 和 dnnlib，或将它们添加到 PYTHONPATH 环境变量中。