StarGAN v2 - 官方 PyTorch 实现

StarGAN v2：多域多样化图像合成
Yunjey Choi*, Youngjung Uh*, Jaejun Yoo*, Jung-Woo Ha
于 CVPR 2020 发表。(* 表示共同第一作者)

论文：https://arxiv.org/abs/1912.01865
视频：https://youtu.be/0EVh5Ki4dIY

摘要： 优秀的图像到图像转换模型应在不同视觉域之间学习映射关系，同时满足以下特性：1) 生成图像的多样性；2) 对多个域的可扩展性。现有方法往往只解决其中一个问题，要么生成的图像缺乏多样性，要么需要为每个域单独训练一个模型。我们提出了 StarGAN v2，这是一个能够同时解决这两个问题的统一框架，并在基准测试中显著优于现有方法。我们在 CelebA-HQ 数据集以及一个新的动物人脸数据集 (AFHQ) 上进行的实验验证了我们在视觉质量、多样性和可扩展性方面的优越性。为了更好地评估图像到图像转换模型，我们发布了 AFHQ 数据集，该数据集包含高质量的动物人脸图像，具有较大的域间和域内差异。代码、预训练模型和数据集均可在 clovaai/stargan-v2 中获取。

预告视频

点击图片观看预告视频。

TensorFlow 实现

我们团队成员 junho 提供的 StarGAN v2 的 TensorFlow 实现可在 clovaai/stargan-v2-tensorflow 找到。

软件安装

克隆本仓库：

git clone https://github.com/clovaai/stargan-v2.git
cd stargan-v2/

安装依赖项：

conda create -n stargan-v2 python=3.6.7
conda activate stargan-v2
conda install -y pytorch=1.4.0 torchvision=0.5.0 cudatoolkit=10.0 -c pytorch
conda install x264=='1!152.20180717' ffmpeg=4.0.2 -c conda-forge
pip install opencv-python==4.1.2.30 ffmpeg-python==0.2.0 scikit-image==0.16.2
pip install pillow==7.0.0 scipy==1.2.1 tqdm==4.43.0 munch==2.5.0

数据集与预训练网络

我们提供了一个脚本，用于下载 StarGAN v2 中使用的数据集及其对应的预训练网络。数据集和网络检查点将分别下载并存储在 data 和 expr/checkpoints 目录中。

CelebA-HQ。 若要下载 CelebA-HQ 数据集及预训练网络，请运行以下命令：

bash download.sh celeba-hq-dataset
bash download.sh pretrained-network-celeba-hq
bash download.sh wing

AFHQ。 若要下载 AFHQ 数据集及预训练网络，请运行以下命令：

bash download.sh afhq-dataset
bash download.sh pretrained-network-afhq

生成插值视频

下载预训练网络后，您可以合成反映参考图像多样化风格（例如发型）的输出图像。以下命令会将生成的图像和插值视频保存到 expr/results 目录中。

CelebA-HQ。 若要生成图像和插值视频，请运行以下命令：

python main.py --mode sample --num_domains 2 --resume_iter 100000 --w_hpf 1 \
               --checkpoint_dir expr/checkpoints/celeba_hq \
               --result_dir expr/results/celeba_hq \
               --src_dir assets/representative/celeba_hq/src \
               --ref_dir assets/representative/celeba_hq/ref

若要转换自定义图像，首先需手动裁剪图像，使面部占整个图像的比例与 CelebA-HQ 数据集相似。然后，运行以下命令进行进一步的微调旋转和裁剪。所有位于 inp_dir 目录中的自定义图像都将被对齐并存储在 out_dir 目录中。

python main.py --mode align \
               --inp_dir assets/representative/custom/female \
               --out_dir assets/representative/celeba_hq/src/female

AFHQ。 若要生成图像和插值视频，请运行以下命令：

python main.py --mode sample --num_domains 3 --resume_iter 100000 --w_hpf 0 \
               --checkpoint_dir expr/checkpoints/afhq \
               --result_dir expr/results/afhq \
               --src_dir assets/representative/afhq/src \
               --ref_dir assets/representative/afhq/ref

评估指标

要使用 Fréchet Inception Distance (FID) 和 Learned Perceptual Image Patch Similarity (LPIPS) 评估 StarGAN v2，请运行以下命令：

# celeba-hq
python main.py --mode eval --num_domains 2 --w_hpf 1 \
               --resume_iter 100000 \
               --train_img_dir data/celeba_hq/train \
               --val_img_dir data/celeba_hq/val \
               --checkpoint_dir expr/checkpoints/celeba_hq \
               --eval_dir expr/eval/celeba_hq

# afhq
python main.py --mode eval --num_domains 3 --w_hpf 0 \
               --resume_iter 100000 \
               --train_img_dir data/afhq/train \
               --val_img_dir data/afhq/val \
               --checkpoint_dir expr/checkpoints/afhq \
               --eval_dir expr/eval/afhq

请注意，评估指标是通过随机潜在向量或参考图像计算得出的，而这些向量或图像均由 seed number 决定。在论文中，我们报告了使用不同 seed number 进行 10 次测量后的平均值。下表展示了基于潜在向量引导和基于参考图像引导的合成所计算出的数值。

数据集	FID (潜在)	LPIPS (潜在)	FID (参考)	LPIPS (参考)	耗时
celeba-hq	13.73 ± 0.06	0.4515 ± 0.0006	23.84 ± 0.03	0.3880 ± 0.0001	49分51秒
afhq	16.18 ± 0.15	0.4501 ± 0.0007	19.78 ± 0.01	0.4315 ± 0.0002	64分49秒

训练网络

要从头开始训练 StarGAN v2，请运行以下命令。生成的图像和网络检查点将分别存储在 expr/samples 和 expr/checkpoints 目录中。在单个 Tesla V100 GPU 上，训练大约需要三天时间。有关训练参数及其说明，请参阅 此处。

# celeba-hq
python main.py --mode train --num_domains 2 --w_hpf 1 \
               --lambda_reg 1 --lambda_sty 1 --lambda_ds 1 --lambda_cyc 1 \
               --train_img_dir data/celeba_hq/train \
               --val_img_dir data/celeba_hq/val

# afhq
python main.py --mode train --num_domains 3 --w_hpf 0 \
               --lambda_reg 1 --lambda_sty 1 --lambda_ds 2 --lambda_cyc 1 \
               --train_img_dir data/afhq/train \
               --val_img_dir data/afhq/val

Animal Faces-HQ 数据集 (AFHQ)

我们发布了一个新的动物人脸数据集——Animal Faces-HQ (AFHQ)，包含 15,000 张 512×512 分辨率的高质量图像。上图展示了 AFHQ 数据集的一些示例图像。该数据集包含猫、狗和野生动物三个类别，每个类别约有 5,000 张图像。由于具有多个（三个）类别，并且每个类别中还包含多种不同的品种，AFHQ 构成了一个极具挑战性的图像到图像翻译问题。对于每个类别，我们选取 500 张图像作为测试集，其余所有图像则作为训练集。要下载该数据集，请运行以下命令：

bash download.sh afhq-dataset

[更新：2021年7月1日] 我们使用高质量的缩放滤波器（即 Lanczos 重采样）重新构建了原始 AFHQ 数据集。请参阅 clean FID 论文，其中指出了下采样过程中不幸存在的软件库问题。我们感谢 Alias-Free GAN 的作者们对更新后的 AFHQ 数据集提出的建议和贡献。如果您使用更新后的数据集，我们建议您不仅引用我们的论文，也引用他们的论文。

与原始数据集相比，更新后的数据集有以下不同之处：

• 我们使用 Lanczos 重采样来调整图像大小，而不是最近邻下采样。
• 原始数据集中约有 2% 的图像已被移除。因此，现在的数据集共有 15,803 张图像，而原始数据集则有 16,130 张。
• 图像以 PNG 格式保存，以避免压缩伪影。这使得文件体积比原始版本更大，但这样做是值得的。

要下载更新后的数据集，请运行以下命令：

bash download.sh afhq-v2-dataset

许可证

源代码、预训练模型和数据集均采用 NAVER Corporation 提供的 知识共享署名-非商业性使用 4.0 许可证。只要您通过引用我们的论文给予 适当署名，并注明是否进行了修改，您就可以出于 非商业目的 对这些材料进行 使用、复制、改编和二次创作。

如需商务合作，请联系 clova-jobs@navercorp.com。
如需技术或其他咨询，请联系 yunjey.choi@navercorp.com。

引用

如果您发现本研究对您的工作有所帮助，请引用我们的论文：

@inproceedings{choi2020starganv2,
  title={StarGAN v2: Diverse Image Synthesis for Multiple Domains},
  author={Yunjey Choi and Youngjung Uh and Jaejun Yoo and Jung-Woo Ha},
  booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
  year={2020}
}

致谢

我们衷心感谢 Clova AI Research（现为 NAVER AI Lab）的全职及访问研究人员提供的宝贵反馈和早期审阅，尤其是 Seongjoon Oh、Junsuk Choe、Muhammad Ferjad Naeem 和 Kyungjune Baek。同时，我们也感谢 Alias-Free GAN 的作者们对更新后的 AFHQ 数据集所做出的贡献。

StarGAN v2 快速上手指南

StarGAN v2 是一个用于多域图像合成的高质量生成模型，支持在单个框架内实现多样化的风格迁移（如人脸属性编辑、动物品种转换等）。本指南帮助中国开发者快速完成环境搭建与基础推理。

环境准备

• 操作系统：Linux（推荐 Ubuntu 18.04+）
• Python 版本：3.6.7
• GPU 要求：支持 CUDA 10.0 的 NVIDIA GPU（训练建议 Tesla V100 或同等性能）
• 前置依赖：
  • Conda（推荐使用 清华镜像源 加速安装）
  • Git

💡 提示：国内用户可配置 Conda 使用清华或中科大镜像源以提升包下载速度：

conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/
conda config --set show_channel_urls yes

安装步骤

1. 克隆项目仓库：

git clone https://github.com/clovaai/stargan-v2.git
cd stargan-v2/

2. 创建并激活 Conda 环境：

conda create -n stargan-v2 python=3.6.7
conda activate stargan-v2

3. 安装核心依赖（PyTorch 及相关工具）：

conda install -y pytorch=1.4.0 torchvision=0.5.0 cudatoolkit=10.0 -c pytorch
conda install x264=='1!152.20180717' ffmpeg=4.0.2 -c conda-forge
pip install opencv-python==4.1.2.30 ffmpeg-python==0.2.0 scikit-image==0.16.2
pip install pillow==7.0.0 scipy==1.2.1 tqdm==4.43.0 munch==2.5.0

⚠️ 若 conda-forge 下载缓慢，可尝试替换为国内镜像源或使用 pip 安装对应包。

基本使用

以下示例演示如何使用预训练模型对 CelebA-HQ 数据集进行风格迁移并生成插值视频。

1. 下载预训练模型与示例数据

bash download.sh celeba-hq-dataset
bash download.sh pretrained-network-celeba-hq
bash download.sh wing

2. 生成风格迁移结果与插值视频

python main.py --mode sample --num_domains 2 --resume_iter 100000 --w_hpf 1 \
               --checkpoint_dir expr/checkpoints/celeba_hq \
               --result_dir expr/results/celeba_hq \
               --src_dir assets/representative/celeba_hq/src \
               --ref_dir assets/representative/celeba_hq/ref

生成的图像和视频将保存在 expr/results/celeba_hq 目录中。

（可选）处理自定义人脸图像

若需使用自己的人脸图片，请先手动裁剪使面部占比接近 CelebA-HQ 风格，然后运行对齐命令：

python main.py --mode align \
               --inp_dir assets/representative/custom/female \
               --out_dir assets/representative/celeba_hq/src/female

处理后的图像将保存至指定输出目录，随后可用于上述 sample 模式进行风格迁移。

---

完成以上步骤后，即可体验 StarGAN v2 的多域多样化图像生成功能。更多高级用法（如训练新模型、评估指标计算等）请参考官方文档。