[项目] [PDF] [补充材料] [演讲] [CGP 摘要] [复现演示] [Hugging Face Spaces 演示]

快速入门，请尝试 Colab： 

本仓库包含我们 NeurIPS 2021 论文《Projected GANs Converge Faster》的代码。

作者：Axel Sauer、Kashyap Chitta、Jens Müller 和 Andreas Geiger。

如果您觉得我们的代码或论文有用，请引用以下文献：

@InProceedings{Sauer2021NEURIPS,
  author         = {Axel Sauer and Kashyap Chitta and Jens M{\"{u}}ller and Andreas Geiger},
  title          = {Projected GANs Converge Faster},
  booktitle      = {Advances in Neural Information Processing Systems (NeurIPS)},
  year           = {2021},
}

相关项目

• StyleGAN-XL: 将 StyleGAN 扩展到大型多样化数据集

待办事项

• 初始代码发布
• 易用的 Colab 笔记本
• 支持 StyleGAN3（已移至 https://github.com/autonomousvision/stylegan_xl）
• 提供预训练模型

环境要求

• 64位 Python 3.8 和 PyTorch 1.9.0（或更高版本）。PyTorch 的安装说明请参见 https://pytorch.org。
• 使用 Miniconda3 创建并激活 PG Python 环境，命令如下：
  • conda env create -f environment.yml
  • conda activate pg
• StyleGAN2 生成器依赖于自定义 CUDA 内核，这些内核会在运行时编译。因此您需要：
  • CUDA 工具包 11.1 或更高版本。
  • GCC 7 或更高版本的编译器。推荐的 GCC 版本取决于 CUDA 版本，例如 CUDA 11.4 的系统要求。
  • 如果在设置自定义 CUDA 内核时遇到问题，您可以参考原始 StyleGAN 仓库的故障排除文档。使用 FastGAN 生成器时，则无需自定义内核。

数据准备

为了快速上手，您可以下载 FastGAN 作者提供的少量样本数据集。您可以在这里下载。要将数据集准备为相应分辨率，可以运行如下命令：

python dataset_tool.py --source=./data/pokemon --dest=./data/pokemon256.zip \
  --resolution=256x256 --transform=center-crop

我们论文中使用的数据集可在其各自官网获取：

CLEVR、FFHQ、Cityscapes、LSUN、AFHQ、Landscape。

训练

使用 8 张 GPU 卡在 LSUN church 数据集上训练您自己的 PG：

python train.py --outdir=./training-runs/ --cfg=fastgan --data=./data/pokemon256.zip \
  --gpus=8 --batch=64 --mirror=1 --snap=50 --batch-gpu=8 --kimg=10000

--batch 指定总批量大小，--batch-gpu 指定每张 GPU 的批量大小。如果您使用的 GPU 数量较少，训练循环会自动累积梯度，直到达到总批量大小。

如果想使用 StyleGAN2 生成器，只需添加 --cfg=stylegan2 参数即可。

我们还提供了一个轻量级的 FastGAN 版本（--cfg=fastgan_lite）。该架构在实际时间上训练速度更快，并且在像 Pokemon 这样的小数据集上表现更好。

样本和指标会保存在 outdir 中。您可以查看 fid50k_full.json 文件，或在 training-runs 目录下运行 tensorboard 来监控训练进度。

生成样本与插值

要生成样本和插值视频，可以运行：

python gen_images.py --outdir=out --trunc=1.0 --seeds=10-15 \
  --network=PATH_TO_NETWORK_PKL

以及

python gen_video.py --output=lerp.mp4 --trunc=1.0 --seeds=0-31 --grid=4x2 \
  --network=PATH_TO_NETWORK_PKL

我们提供了以下预训练模型（将 URL 替换为 PATH_TO_NETWORK_PKL）：

_{https://s3.eu-central-1.amazonaws.com/avg-projects/projected_gan/models/art_painting.pkl
https://s3.eu-central-1.amazonaws.com/avg-projects/projected_gan/models/church.pkl
https://s3.eu-central-1.amazonaws.com/avg-projects/projected_gan/models/bedroom.pkl
https://s3.eu-central-1.amazonaws.com/avg-projects/projected_gan/models/cityscapes.pkl
https://s3.eu-central-1.amazonaws.com/avg-projects/projected_gan/models/clevr.pkl
https://s3.eu-central-1.amazonaws.com/avg-projects/projected_gan/models/ffhq.pkl
https://s3.eu-central-1.amazonaws.com/avg-projects/projected_gan/models/flowers.pkl
https://s3.eu-central-1.amazonaws.com/avg-projects/projected_gan/models/landscape.pkl
https://s3.eu-central-1.amazonaws.com/avg-projects/projected_gan/models/pokemon.pkl}

质量指标

默认情况下，train.py 会在训练过程中跟踪 FID50k 指标。要计算特定网络快照的指标，可以运行：

python calc_metrics.py --metrics=fid50k_full --network=PATH_TO_NETWORK_PKL

要查看可用的指标，可以运行：

python calc_metrics.py --help

在您自己的项目中使用 PG

我们的实现是模块化的，因此很容易将其集成到您自己的代码库中。只需将 pg_modules 文件夹复制到您的项目中即可。

然后，要获取投影式多尺度判别器，可以运行：

from pg_modules.discriminator import ProjectedDiscriminator
D = ProjectedDiscriminator()

您唯一需要做的就是确保特征网络不参与训练，即在训练循环中显式设置：

D.feature_network.requires_grad_(False)

致谢

我们的代码基于并扩展了 Karras 等人开发的优秀 StyleGAN2-ADA 仓库 和 StyleGAN3 仓库。

此外，我们还使用了 FastGAN 和 MiDas 的部分代码。

Projected GAN 快速上手指南

Projected GAN 是一种收敛速度更快的生成对抗网络，特别适合小数据集训练。本指南将帮助您快速搭建环境并运行模型。

环境准备

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

• 操作系统: 64-bit Linux (推荐) 或 Windows
• Python: 3.8 或更高版本
• PyTorch: 1.9.0 或更高版本
• GPU: NVIDIA 显卡，需安装 CUDA Toolkit 11.1 或更高版本
• 编译器: GCC 7 或更高版本（用于编译自定义 CUDA 内核）

注意: 如果您计划使用 FastGAN 生成器而非 StyleGAN2 生成器，则不需要配置复杂的 CUDA 编译环境。

安装步骤

推荐使用 Miniconda3 创建独立的虚拟环境以避免依赖冲突。

1. 克隆仓库

   git clone https://github.com/autonomousvision/projected_gan.git
   cd projected_gan
   

2. 创建并激活 Conda 环境 使用项目提供的配置文件一键安装依赖：

   conda env create -f environment.yml
   conda activate pg
   

3. 验证 CUDA 环境 确保 nvcc --version 和 gcc --version 输出版本符合要求。如果遇到自定义 CUDA 内核编译错误，请参考 StyleGAN3 的故障排除文档。

基本使用

1. 准备数据

您可以下载作者提供的少样本数据集（如 Pokemon）进行快速测试，或使用自己的图片文件夹。

示例：处理 Pokemon 数据集为 256x256 分辨率

python dataset_tool.py --source=./data/pokemon --dest=./data/pokemon256.zip \
  --resolution=256x256 --transform=center-crop

2. 开始训练

以下命令演示如何在 8 张 GPU 上使用 fastgan_lite 配置（轻量级，适合小数据集）训练模型：

python train.py --outdir=./training-runs/ --cfg=fastgan_lite --data=./data/pokemon256.zip \
  --gpus=8 --batch=64 --mirror=1 --snap=50 --batch-gpu=8 --kimg=10000

• --cfg: 可选 fastgan, fastgan_lite, 或 stylegan2。
• --batch: 总批次大小。
• --batch-gpu: 单卡批次大小（若显卡较少，程序会自动累积梯度）。
• 训练进度和样本保存在 --outdir 指定目录，可使用 TensorBoard 监控。

3. 生成图像与视频

使用预训练模型或训练好的快照生成样本。

生成静态图像：

python gen_images.py --outdir=out --trunc=1.0 --seeds=10-15 \
  --network=https://s3.eu-central-1.amazonaws.com/avg-projects/projected_gan/models/pokemon.pkl

生成插值视频：

python gen_video.py --output=lerp.mp4 --trunc=1.0 --seeds=0-31 --grid=4x2 \
  --network=https://s3.eu-central-1.amazonaws.com/avg-projects/projected_gan/models/pokemon.pkl

预训练模型链接: 项目提供了多个领域的预训练模型，可直接作为 --network 参数传入：

• Art Painting: .../models/art_painting.pkl
• Church: .../models/church.pkl
• Bedroom: .../models/bedroom.pkl
• Cityscapes: .../models/cityscapes.pkl
• CLEVR: .../models/clevr.pkl
• FFHQ: .../models/ffhq.pkl
• Flowers: .../models/flowers.pkl
• Landscape: .../models/landscape.pkl
• Pokemon: .../models/pokemon.pkl

4. 在自有项目中集成

Projected GAN 采用模块化设计，只需将 pg_modules 文件夹复制到您的项目中即可调用判别器：

from pg_modules.discriminator import ProjectedDiscriminator

# 初始化投影多尺度判别器
D = ProjectedDiscriminator()

# 重要：确保特征网络不参与训练
D.feature_network.requires_grad_(False)