PyTorch 中的 StyleGAN 2

在 PyTorch 中实现对 StyleGAN 图像质量的分析与改进（https://arxiv.org/abs/1912.04958）

注意事项

我已尽力使此实现尽可能接近官方实现，但可能仍有一些细节未被注意到。因此，请谨慎使用此实现。

要求

我在以下环境中进行了测试：

• PyTorch 1.3.1
• CUDA 10.1/10.2

使用方法

首先创建 LMDB 数据集：

python prepare_data.py --out LMDB_PATH --n_worker N_WORKER --size SIZE1,SIZE2,SIZE3,... DATASET_PATH

此命令会将图像转换为 JPEG 格式并预先调整其大小。该实现并未采用渐进式生成技术，但您可以通过逗号分隔的分辨率列表来创建多个分辨率的数据集，以便日后尝试其他分辨率。

然后，您可以在分布式环境中训练模型：

python -m torch.distributed.launch --nproc_per_node=N_GPU --master_port=PORT train.py --batch BATCH_SIZE LMDB_PATH

train.py 支持 Weights & Biases 日志记录。如果您希望使用此功能，请在脚本中添加 --wandb 参数。

SWAGAN

此实现实验性地支持 SWAGAN：一种基于风格的小波驱动生成模型（https://arxiv.org/abs/2102.06108）。您可以通过以下命令训练 SWAGAN：

python -m torch.distributed.launch --nproc_per_node=N_GPU --master_port=PORT train.py --arch swagan --batch BATCH_SIZE LMDB_PATH

正如论文中所指出的那样，SWAGAN 的训练速度要快得多（在 256 像素分辨率下大约快 2 倍）。

从官方检查点转换权重

您需要克隆官方仓库（https://github.com/NVlabs/stylegan2），因为加载官方检查点需要这些代码。

例如，如果您将仓库克隆到 ~/stylegan2 目录，并下载了 stylegan2-ffhq-config-f.pkl 文件，您可以按如下方式进行转换：

python convert_weight.py --repo ~/stylegan2 stylegan2-ffhq-config-f.pkl

这将生成一个转换后的 stylegan2-ffhq-config-f.pt 文件。

生成样本

python generate.py --sample N_FACES --pics N_PICS --ckpt PATH_CHECKPOINT

如果您使用其他尺寸进行训练，则应相应地更改尺寸参数（例如，设置为 256）。

将图像投影到潜在空间

python projector.py --ckpt [CHECKPOINT] --size [GENERATOR_OUTPUT_SIZE] FILE1 FILE2 ...

闭式分解（https://arxiv.org/abs/2007.06600）

您可以使用 closed_form_factorization.py 和 apply_factor.py 来无监督地发现有意义的潜在语义因子或方向。

首先，您需要使用 closed_form_factorization.py 提取权重矩阵的特征向量：

python closed_form_factorization.py [CHECKPOINT]

这将生成包含特征向量的文件（默认为 factor.pt）。然后，您可以使用 apply_factor.py 来测试提取的方向的意义：

python apply_factor.py -i [INDEX_OF_EIGENVECTOR] -d [DEGREE_OF_MOVE] -n [NUMBER_OF_SAMPLES] --ckpt [CHECKPOINT] [FACTOR_FILE]

例如：

python apply_factor.py -i 19 -d 5 -n 10 --ckpt [CHECKPOINT] factor.pt

这将生成 10 个随机样本，这些样本由沿着第 19 个特征向量移动、幅度为 ±5 的潜在变量生成。

预训练检查点

链接

我在 FFHQ 数据集上训练了一个 256 像素的模型，共进行了 55 万次迭代。得到的 FID 约为 4.5。可能是数据预处理、分辨率或训练循环等因素导致了这一差异，但我目前尚不清楚 FID 差异的确切原因。

样本

来自 FFHQ 的样本，训练至 11 万次迭代（基于 352 万张图像）。

来自 MetFaces 的样本，使用了非泄漏增强技术，训练至 15 万次迭代（基于 480 万张图像）。

来自转换后权重的样本

FFHQ 数据集上的 1024 像素样本。

LSUN 教堂数据集上的 256 像素样本。

许可证

模型细节和自定义 CUDA 内核代码均来自官方仓库：https://github.com/NVlabs/stylegan2

用于学习感知图像块相似度（LPIPS）的代码来自 https://github.com/richzhang/PerceptualSimilarity

为了使 FID 分数更接近 TensorFlow 官方实现，我使用了 https://github.com/mseitzer/pytorch-fid 中的 FID Inception V3 实现。

StyleGAN2-PyTorch 快速上手指南

本指南基于 stylegan2-pytorch 项目，帮助开发者快速在 PyTorch 环境中部署和运行 StyleGAN 2 模型。

环境准备

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

• 操作系统: Linux (推荐)
• Python: 3.x
• PyTorch: 1.3.1 或更高版本
• CUDA: 10.1 或 10.2 (需与 PyTorch 版本匹配)
• 依赖库:
  • torchvision
  • ninja (用于编译自定义 CUDA 内核)
  • lmdb
  • wandb (可选，用于训练日志记录)

提示：国内用户建议使用清华源或阿里源加速 PyTorch 及相关依赖的安装。

pip install torch torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install ninja lmdb wandb -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

1. 克隆仓库

   git clone https://github.com/rosinality/stylegan2-pytorch.git
   cd stylegan2-pytorch
   

2. 验证环境 确保 nvcc --version 显示的 CUDA 版本与安装的 PyTorch 版本兼容。首次运行脚本时，项目会自动编译自定义 CUDA 算子。

基本使用

1. 准备数据集 (LMDB 格式)

StyleGAN2-PyTorch 需要将图像数据集转换为 LMDB 格式以提升读取效率。假设您有一个包含图像的文件夹 DATASET_PATH：

python prepare_data.py --out LMDB_PATH --n_worker 4 --size 256 DATASET_PATH

• --out: 输出 LMDB 文件的路径。
• --n_worker: 并行处理的工作线程数。
• --size: 生成图像的尺寸（支持逗号分隔的多分辨率，如 256,512,1024）。
• 该命令会将图像转换为 JPEG 并预调整大小。

2. 训练模型

支持分布式训练。以下命令示例使用 2 张 GPU 进行训练：

python -m torch.distributed.launch --nproc_per_node=2 --master_port=29500 train.py --batch 8 LMDB_PATH

• --nproc_per_node: 使用的 GPU 数量。
• --batch: 总批次大小（会自动根据 GPU 数量分配）。
• 可选: 添加 --wandb 参数以启用 Weights & Biases 日志记录。

实验性功能 (SWAGAN): 如需训练速度更快的 SWAGAN 模型，添加 --arch swagan 参数：

python -m torch.distributed.launch --nproc_per_node=2 --master_port=29500 train.py --arch swagan --batch 8 LMDB_PATH

3. 加载官方预训练权重 (可选)

如果您想直接使用 NVIDIA 官方发布的权重，需要先克隆官方仓库进行格式转换：

# 1. 克隆官方仓库 (用于转换脚本依赖)
git clone https://github.com/NVlabs/stylegan2.git ~/stylegan2

# 2. 转换权重 (假设已下载 stylegan2-ffhq-config-f.pkl)
python convert_weight.py --repo ~/stylegan2 stylegan2-ffhq-config-f.pkl

转换后将生成 stylegan2-ffhq-config-f.pt 文件。

4. 生成样本

使用训练好的模型或转换后的官方权重生成图像：

python generate.py --sample 10 --pics 5 --ckpt path/to/checkpoint.pt --size 256

• --sample: 每次生成的样本数量。
• --pics: 生成的图片总数。
• --ckpt: 模型检查点路径 (.pt 文件)。
• --size: 生成图像尺寸，需与训练时的分辨率一致。

5. 潜在空间探索 (进阶)

• 图像投影: 将真实图像投影到潜在空间。

  python projector.py --ckpt [CHECKPOINT] --size 256 image1.jpg image2.jpg
  

• 闭式因子分解: 发现潜在的语义方向。

  # 提取特征向量
  python closed_form_factorization.py [CHECKPOINT]
  
  # 应用因子生成变化样本
  python apply_factor.py -i 19 -d 5 -n 10 --ckpt [CHECKPOINT] factor.pt