对比无配对图像翻译 (CUT)

视频 (1分钟) | 视频 (10分钟) | 网站 | 论文

我们提供基于局部对比学习和对抗学习的无配对图像到图像翻译的 PyTorch 实现。该方法未使用手工设计的损失函数或逆向网络。与 CycleGAN 相比，我们的模型训练速度更快且内存占用更少。此外，我们的方法还可以扩展到单张图像训练，其中每个“域”仅由一张图像构成。

用于无配对图像到图像翻译的对比学习
Taesung Park、Alexei A. Efros、Richard Zhang、Jun-Yan Zhu
加州大学伯克利分校和 Adobe Research
在 ECCV 2020 上发表

伪代码

import torch
cross_entropy_loss = torch.nn.CrossEntropyLoss()

# 输入：f_q (BxCxS) 和从 H(G_enc(x)) 中采样的特征
# 输入：f_k (BxCxS) 是从 H(G_enc(G(x))) 中采样的特征
# 输入：tau 是 PatchNCE 损失中使用的温度。
# 输出：PatchNCE 损失
def PatchNCELoss(f_q, f_k, tau=0.07):
    # 批量大小、通道大小和采样位置的数量
    B, C, S = f_q.shape

    # 计算 v * v+：BxSx1
    l_pos = (f_k * f_q).sum(dim=1)[:, :, None]

    # 计算 v * v-：BxSxS
    l_neg = torch.bmm(f_q.transpose(1, 2), f_k)

    # 对角线元素不是负样本，将其移除。
    identity_matrix = torch.eye(S)[None, :, :]
    l_neg.masked_fill_(identity_matrix, -float('inf'))

    # 计算 logits：(B)x(S)x(S+1)
    logits = torch.cat((l_pos, l_neg), dim=2) / tau

    # 返回 PatchNCE 损失
    predictions = logits.flatten(0, 1)
    targets = torch.zeros(B * S, dtype=torch.long)
    return cross_entropy_loss(predictions, targets)

示例结果

无配对图像到图像翻译

单张图像无配对翻译

俄罗斯蓝猫到暴躁猫

巴黎街头到布拉诺彩绘房屋

先决条件

• Linux 或 macOS
• Python 3
• CPU 或 NVIDIA GPU + CUDA CuDNN

更新日志

2020年9月12日：增加了单张图像翻译功能。

开始使用

• 克隆此仓库：

git clone https://github.com/taesungp/contrastive-unpaired-translation CUT
cd CUT

• 安装 PyTorch 1.1 及其他依赖项（如 torchvision、visdom、dominate、gputil）。

  对于 pip 用户，请输入命令 pip install -r requirements.txt。

  对于 Conda 用户，可以使用 conda env create -f environment.yml 创建一个新的 Conda 环境。

CUT 和 FastCUT 的训练与测试

• 下载 grumpifycat 数据集（论文图8，俄罗斯蓝猫→暴躁猫）

bash ./datasets/download_cut_dataset.sh grumpifycat

数据集将被下载并解压到 ./datasets/grumpifycat/。

• 要查看训练结果和损失曲线，请运行 python -m visdom.server 并点击 URL http://localhost:8097。

• 训练 CUT 模型：

python train.py --dataroot ./datasets/grumpifycat --name grumpycat_CUT --CUT_mode CUT

或者训练 FastCUT 模型

python train.py --dataroot ./datasets/grumpifycat --name grumpycat_FastCUT --CUT_mode FastCUT

检查点将存储在 ./checkpoints/grumpycat_*/web。

• 测试 CUT 模型：

python test.py --dataroot ./datasets/grumpifycat --name grumpycat_CUT --CUT_mode CUT --phase train

测试结果将保存为 HTML 文件，位于 ./results/grumpifycat/latest_train/index.html。

CUT、FastCUT 和 CycleGAN

CUT 使用身份保持损失并设置 lambda_NCE=1 进行训练，而 FastCUT 则不使用身份损失，但设置了更高的 lambda_NCE=10.0。与 CycleGAN 相比，CUT 学习进行更强大的分布匹配，而 FastCUT 则是一种更轻量级（GPU 内存占用仅为 CycleGAN 的一半，可处理更大尺寸的图像）、训练速度更快（是 CycleGAN 的两倍）的替代方案。更多细节请参阅 论文。

在上图中，我们使用预先训练好的语义分割模型来测量属于马或斑马身体的像素比例。我们发现马和斑马图像的尺寸分布存在差异——斑马通常显得更大（36.8% vs. 17.9%）。我们的完整方法 CUT 具有放大马的能力，从而能够更好地匹配训练数据的统计特性，这比 CycleGAN 更为出色。而 FastCUT 的表现则更保守，类似于 CycleGAN。

使用我们的启动脚本进行训练

请参阅 experiments/grumpifycat_launcher.py，它会生成上述命令行参数。这些启动脚本对于配置复杂度较高的训练和测试命令行参数非常有用。

通过启动脚本，以下命令可以生成 CUT 和 FastCUT 的训练命令。

python -m experiments grumpifycat train 0   # CUT
python -m experiments grumpifycat train 1   # FastCUT

要使用启动脚本进行测试，

python -m experiments grumpifycat test 0   # CUT
python -m experiments grumpifycat test 1   # FastCUT

可用的命令包括 run、run_test、launch、close 等。所有命令请参阅 experiments/__main__.py。启动脚本易于定义和使用。例如，grumpifycat 启动脚本仅用几行代码即可定义：

from .tmux_launcher import Options、TmuxLauncher


class Launcher(TmuxLauncher):
    def common_options(self):
        return [
            Options(    # 命令 0
                dataroot="./datasets/grumpifycat",
                name="grumpifycat_CUT",
                CUT_mode="CUT"
            ),

            Options(    # 命令 1
                dataroot="./datasets/grumpifycat",
                name="grumpifycat_FastCUT",
                CUT_mode="FastCUT",
            )
        ]

    def commands(self):
        return ["python train.py " + str(opt) for opt in self.common_options()]

    def test_commands(self):
        # 俄罗斯蓝猫→暴躁猫的数据集没有测试集划分。
        # 因此，我们将测试集设为“训练集”。
        return ["python test.py " + str(opt.set(phase='train')) for opt in self.common_options()]

应用预训练的 CUT 模型并评估 FID

要运行预训练的模型，请执行以下操作。

# 下载并解压预训练的模型。权重应位于

# 例如，检查点/horse2zebra_cut_pretrained/latest_net_G.pth。
wget http://efrosgans.eecs.berkeley.edu/CUT/pretrained_models.tar
tar -xf pretrained_models.tar

# 生成输出。数据集路径可能需要调整。
# 要做到这一点，修改experiments/pretrained_launcher.py中的相应行。
# [id] 对应于pretrained_launcher.py中定义的各个命令：
# 0 - 在Cityscapes数据集上使用CUT
# 1 - 在Cityscapes数据集上使用FastCUT
# 2 - 在Horse2Zebra数据集上使用CUT
# 3 - 在Horse2Zebra数据集上使用FastCUT
# 4 - 在Cat2Dog数据集上使用CUT
# 5 - 在Cat2Dog数据集上使用FastCUT
python -m experiments pretrained run_test [id]

# 评估FID。为此，首先安装https://github.com/mseitzer/pytorch-fid上的pytorch-fid。
# pip install pytorch-fid
# 例如，要评估Horse2Zebra数据集上CUT模型的FID，
# python -m pytorch_fid ./datasets/horse2zebra/testB/ results/horse2zebra_cut_pretrained/test_latest/images/fake_B/
# 要评估Cityscapes数据集上FastCUT模型的FID，
# python -m pytorch_fid ./datasets/cityscapes/valA/ ~/projects/contrastive-unpaired-translation/results/cityscapes_fastcut_pretrained/test_latest/images/fake_B/
# 注意，Cityscapes模型需要使用特殊的数据集。请参阅下文。
python -m pytorch_fid [真实测试图像路径] [生成图像路径]

注意：Cityscapes预训练模型是在原始Cityscapes数据集的缩放和JPEG压缩版本上进行训练和评估的。要进行评估，请下载此链接中的验证集并进行评估。

SinCUT 单张图像非配对训练

要训练SinCUT（单张图像翻译，如论文中的图9、13和14所示），您需要：

1. 将--model选项设置为--model sincut，这将调用位于./models/sincut_model.py中的配置和代码；以及
2. 指定每个域中包含一张图像的数据集目录，例如本仓库中包含的示例数据集./datasets/single_image_monet_etretat/。

例如，要训练一个用于埃特雷塔悬崖（图13的第一张图片）的模型，请使用以下命令。

python train.py --model sincut --name singleimage_monet_etretat --dataroot ./datasets/single_image_monet_etretat

或者使用实验启动脚本：

python -m experiments singleimage run 0

对于单张图像翻译，我们采用了StyleGAN2的网络架构组件，以及在DTN和CycleGAN中使用的像素身份保持损失。特别是，我们借鉴了rosinality的代码，该代码位于models/stylegan_networks.py中。

训练过程需要几个小时。要使用检查点生成最终图像，

python test.py --model sincut --name singleimage_monet_etretat --dataroot ./datasets/single_image_monet_etretat

或者简单地：

python -m experiments singleimage run_test 0

数据集

下载CUT/CycleGAN/pix2pix数据集。例如，

bash datasets/download_cut_datasets.sh horse2zebra

Cat2Dog数据集是从AFHQ数据集准备的。请访问https://github.com/clovaai/stargan-v2，通过该仓库的`bash download.sh afhq-dataset`下载AFHQ数据集。然后按如下方式重新组织目录：

mkdir datasets/cat2dog
ln -s datasets/cat2dog/trainA [path_to_afhq]/train/cat
ln -s datasets/cat2dog/trainB [path_to_afhq]/train/dog
ln -s datasets/cat2dog/testA [path_to_afhq]/test/cat
ln -s datasets/cat2dog/testB [path_to_afhq]/test/dog

Cityscapes数据集可以从https://cityscapes-dataset.com下载。之后，使用`./datasets/prepare_cityscapes_dataset.py`脚本来准备数据集。

输入图像的预处理

输入图像的预处理，如调整大小或随机裁剪，由--preprocess、--load_size和--crop_size选项控制。其用法遵循CycleGAN/pix2pix仓库中的做法。

例如，默认设置--preprocess resize_and_crop --load_size 286 --crop_size 256会将输入图像调整为286x286大小，然后随机裁剪出256x256的区域，以此进行数据增强。还有其他可指定的预处理选项，这些选项在base_dataset.py中定义。以下是一些示例选项：

• --preprocess none：不进行任何预处理。请注意，图像尺寸仍会被缩放到最接近的4的倍数，因为卷积生成器无法在其他情况下保持相同的图像尺寸。
• --preprocess scale_width --load_size 768：将图像宽度缩放到768。
• --preprocess scale_shortside_and_crop：按比例缩放图像，使短边达到load_size大小，然后随机裁剪出crop_size大小的窗口。

可以通过修改base_dataset.py中的get_transform()来添加更多预处理选项。

引用

如果您在研究中使用此代码，请引用我们的论文。

@inproceedings{park2020cut,
  title={Contrastive Learning for Unpaired Image-to-Image Translation},
  author={Taesung Park and Alexei A. Efros and Richard Zhang and Jun-Yan Zhu},
  booktitle={European Conference on Computer Vision},
  year={2020}
}

如果您使用本仓库中包含的原始pix2pix和CycleGAN模型，请引用以下论文：

@inproceedings{CycleGAN2017,
  title={Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks},
  author={Zhu, Jun-Yan and Park, Taesung and Isola, Phillip and Efros, Alexei A},
  booktitle={IEEE International Conference on Computer Vision (ICCV)},
  year={2017}
}


@inproceedings{isola2017image,
  title={Image-to-Image Translation with Conditional Adversarial Networks},
  author={Isola, Phillip and Zhu, Jun-Yan and Zhou, Tinghui and Efros, Alexei A},
  booktitle={IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
  year={2017}
}

致谢

我们感谢艾伦·贾布里和菲利普·伊索拉提供的有益讨论和反馈。我们的代码基于 pytorch-CycleGAN-and-pix2pix 开发。同时，我们也感谢 pytorch-fid 提供的 FID 计算功能、drn 提供的 mIoU 计算功能，以及 stylegan2-pytorch 为我们单图像转换任务中使用的 StyleGAN2 的 PyTorch 实现所提供的支持。

Contrastive Unpaired Translation (CUT) 快速上手指南

CUT 是一个基于对比学习的非配对图像到图像翻译工具。相比 CycleGAN，它具有训练更快、显存占用更低的特点，并支持单张图像训练模式（SinCUT）。

环境准备

• 操作系统: Linux 或 macOS
• Python: Python 3
• 硬件: CPU 或 NVIDIA GPU (需安装 CUDA 和 CuDNN)
• 深度学习框架: PyTorch 1.1+

国内加速建议： 在安装依赖时，推荐使用清华或阿里镜像源以加快下载速度。

• pip: pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
• conda: 配置 .condarc 使用清华源 (https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/)

安装步骤

1. 克隆仓库

   git clone https://github.com/taesungp/contrastive-unpaired-translation CUT
   cd CUT
   

2. 安装依赖 选择以下任一方式安装所需库（如 torchvision, visdom, dominate 等）：

   • 使用 pip:

     pip install -r requirements.txt
     # 国内用户推荐：
     # pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
     

   • 使用 Conda:

     conda env create -f environment.yml
     conda activate cut
     

基本使用

以下以官方提供的 "grumpifycat" (蓝猫变不爽猫) 数据集为例，演示从数据下载到训练、测试的全流程。

1. 下载示例数据集

bash ./datasets/download_cut_dataset.sh grumpifycat

数据将解压至 ./datasets/grumpifycat/ 目录。

2. 启动可视化监控 (可选)

在终端运行以下命令，然后在浏览器访问 http://localhost:8097 查看训练损失和结果图。

python -m visdom.server

3. 训练模型

你可以选择训练标准 CUT 模型或更轻量的 FastCUT 模型。

• 训练 CUT 模型:

  python train.py --dataroot ./datasets/grumpifycat --name grumpycat_CUT --CUT_mode CUT
  

• 训练 FastCUT 模型 (速度更快，显存占用减半):

  python train.py --dataroot ./datasets/grumpifycat --name grumpycat_FastCUT --CUT_mode FastCUT
  

  检查点将保存在 ./checkpoints/grumpycat_*/web。

4. 测试与生成结果

使用训练好的模型生成转换后的图像。

python test.py --dataroot ./datasets/grumpifycat --name grumpycat_CUT --CUT_mode CUT --phase train

注意：由于该示例数据集没有独立的测试集，因此 --phase 设置为 train。

生成的结果将保存为 HTML 文件，路径如下： ./results/grumpifycat/latest_train/index.html 直接在浏览器打开该文件即可查看对比效果。

---

进阶提示： 若需进行单张图像风格迁移（SinCUT），请使用 --model sincut 参数，并将数据根目录指向仅包含单张图片的文件夹。