pix2pixHD

项目 | YouTube | 论文

这是我们提出的用于高分辨率（例如 2048x1024）逼真图像到图像转换方法的 PyTorch 实现。它可以用于将语义标签图转换为照片级真实的图像，或从人脸标签图合成肖像。

使用条件 GAN 进行高分辨率图像合成与语义操控
王廷淳¹, 刘明宇¹, 朱俊彦², 安德鲁·陶¹, 扬·考茨¹, 布莱恩·卡坦扎罗^1
1NVIDIA 公司, ²加州大学伯克利分校
发表于 CVPR 2018。

2k/1k 分辨率下的图像到图像转换

• 我们的标签到街景结果

- 交互式编辑结果

- 更多的街景结果

• 标签到人脸及交互式编辑结果

• 我们的编辑界面

前置条件

• Linux 或 macOS
• Python 2 或 3
• NVIDIA GPU（显存 11G 或以上）+ CUDA cuDNN

快速入门

安装

• 从 http://pytorch.org 安装 PyTorch 和依赖项
• 安装 Python 库 dominate。

pip install dominate

• 克隆本仓库：

git clone https://github.com/NVIDIA/pix2pixHD
cd pix2pixHD

测试

• datasets 文件夹中包含了几张 Cityscapes 的示例测试图片。
• 请从 这里 下载预训练的 Cityscapes 模型（Google Drive 链接），并将其放置在 ./checkpoints/label2city_1024p/ 目录下。
• 测试模型（bash ./scripts/test_1024p.sh）：

#!./scripts/test_1024p.sh
python test.py --name label2city_1024p --netG local --ngf 32 --resize_or_crop none

测试结果将保存到此处的 HTML 文件：./results/label2city_1024p/test_latest/index.html。

更多示例脚本可在 scripts 目录中找到。

数据集

• 我们使用 Cityscapes 数据集。若要基于完整数据集训练模型，请从 官方网站 下载（需注册）。 下载后，请按照示例图片的方式将其放入 datasets 文件夹中。

训练

• 在 1024 x 512 分辨率下训练模型（bash ./scripts/train_512p.sh）：

#!./scripts/train_512p.sh
python train.py --name label2city_512p

• 要查看训练结果，请访问 ./checkpoints/label2city_512p/web/index.html 中的中间结果。 如果已安装 TensorFlow，可通过在训练脚本中添加 --tf_log 来查看 ./checkpoints/label2city_512p/logs 中的 TensorBoard 日志。

多 GPU 训练

• 使用多块 GPU 训练模型（bash ./scripts/train_512p_multigpu.sh）：

#!./scripts/train_512p_multigpu.sh
python train.py --name label2city_512p --batchSize 8 --gpu_ids 0,1,2,3,4,5,6,7

注意：此功能未经充分测试，我们仅使用单 GPU 训练了模型。请根据自身情况谨慎使用。

使用自动混合精度 (AMP) 加速训练

• 要启用混合精度支持，请先从 https://github.com/NVIDIA/apex 安装 apex。
• 然后可以通过添加 --fp16 来训练模型。例如，

#!./scripts/train_512p_fp16.sh
python -m torch.distributed.launch train.py --name label2city_512p --fp16

在我们的测试中，使用 Volta 架构的机器时，启用 AMP 后训练速度提升了约 80%。

全分辨率训练

• 要以全分辨率（2048 x 1024）训练图像，需要配备 24G 显存的 GPU（bash ./scripts/train_1024p_24G.sh），或在使用混合精度 (AMP) 的情况下配备 16G 显存。
• 如果仅有 12G 显存的 GPU 可用，请使用 12G 版本的脚本（bash ./scripts/train_1024p_12G.sh），该脚本会在训练过程中对图像进行裁剪。但无法保证在此设置下的性能。

使用自定义数据集训练

• 如果您希望使用自己的数据集进行训练，需生成单通道的标签图，其像素值对应于对象标签（即 0,1,...,N-1，其中 N 为标签数量）。这是因为我们需要从标签图中生成独热向量。此外，在训练和测试时还需指定 --label_nc N。
• 如果您的输入不是标签图，则只需指定 --label_nc 0，系统将直接使用 RGB 颜色作为输入。此时文件夹应命名为 train_A、train_B，而非 train_label、train_img，目标是将 A 类图像转换为 B 类图像。
• 如果您没有实例图或不想使用它们，请指定 --no_instance。
• 默认的预处理设置是 scale_width，它会将所有训练图像的宽度缩放到 opt.loadSize（1024），同时保持宽高比不变。如果您希望采用其他设置，可以使用 --resize_or_crop 选项进行调整。例如，scale_width_and_crop 会先将图像调整至 opt.loadSize 宽度，再随机裁剪为 (opt.fineSize, opt.fineSize) 大小。而 crop 则跳过缩放步骤，仅进行随机裁剪。若您不希望进行任何预处理，可指定 none，此时系统仅确保图像尺寸能被 32 整除。

更多训练/测试细节

• 标志位：所有训练标志请参见 options/train_options.py 和 options/base_options.py；所有测试标志请参见 options/test_options.py 和 options/base_options.py。
• 实例图：我们同时接受标签图和实例图作为输入。如果您不想使用实例图，请指定 --no_instance 标志。

引用

如果您觉得本文对您的研究有帮助，请使用以下引用。

@inproceedings{wang2018pix2pixHD,
  title={利用条件生成对抗网络进行高分辨率图像合成与语义操控},
  author={王婷淳、刘明宇、朱俊彦、安德鲁·陶、扬·考茨、布莱恩·卡坦扎罗},  
  booktitle={IEEE计算机视觉与模式识别会议论文集},
  year={2018}
}

致谢

本代码大量借鉴了 pytorch-CycleGAN-and-pix2pix。

pix2pixHD 快速上手指南

pix2pixHD 是一个基于 PyTorch 的高分辨率（如 2048x1024）图像到图像翻译工具，主要用于将语义标签图转换为照片级真实图像，或从人脸标签图合成人像。

环境准备

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

• 操作系统: Linux 或 macOS
• Python 版本: Python 2 或 3
• 硬件要求: NVIDIA GPU (显存 11GB 或以上) + CUDA cuDNN
  • 注：训练全分辨率 (2048x1024) 模型建议显存 24GB，或使用混合精度训练 (AMP) 降至 16GB。
• 核心依赖:
  • PyTorch (请访问 pytorch.org 安装适合您环境的版本)
  • dominate 库

安装步骤

1. 安装 Python 依赖库

   pip install dominate
   

   (国内用户推荐使用清华源加速：pip install dominate -i https://pypi.tuna.tsinghua.edu.cn/simple)

2. 克隆项目代码

   git clone https://github.com/NVIDIA/pix2pixHD
   cd pix2pixHD
   

基本使用

以下是最简单的测试流程，使用预训练的 Cityscapes 模型将标签图转换为街景图像。

1. 下载预训练模型

请下载预训练的 Cityscapes 模型 (Google Drive 链接)，并将其放置在以下目录中： ./checkpoints/label2city_1024p/

(注：如果无法访问 Google Drive，需自行寻找国内镜像源或手动下载后上传至服务器)

2. 运行测试脚本

执行以下命令进行测试（默认使用 datasets 文件夹中的示例图片）：

python test.py --name label2city_1024p --netG local --ngf 32 --resize_or_crop none

3. 查看结果

测试完成后，生成的图像和对比结果将保存在 HTML 文件中，请使用浏览器打开查看： ./results/label2city_1024p/test_latest/index.html

---

如需使用自己的数据集进行训练，请参考原文档 "Training with your own dataset" 章节准备标签图（单通道，像素值对应类别索引），并使用 --label_nc N 参数指定类别数量。