UVCGAN：UNet Vision Transformer cycle-consistent GAN for unpaired image-to-image translation

[2022ICCV] UVCGAN：UNet Vision Transformer cycle-consistent GAN for unpaired image-to-image translation

文章：https://arxiv.org/pdf/2203.02557v3.pdf

代码：https://github.com/LS4GAN/uvcgan

摘要：

更近期的工作促进一对多映射，以提高翻译图像的多样性。在科学模拟和一对一需求的推动下，这项工作重新审视了经典的CycleGAN框架，并在不放松周期一致性约束的情况下提高了其性能，使其优于更现代的模型。为了实现这一点，

我们为CycleGAN生成器配备Vision Transformer(ViT)来提高非局部模式学习和采用先进的训练技术，如梯度惩罚和自我监督预训练，得到的模型，称为UVCGAN，可以在几个基准数据集中优于竞争模型。梯度惩罚和自我监督的预训练是改进的关键。

为了促进更多样化的图像生成，许多最近的作品[80,56,61,79]放松了循环一致性约束。

56：Breaking the cycle–colleagues are all you need. （Council-GAN）

61：Contrastive learning for unpaired image-to-image translation.（CUT）

79：Unpaired image-to-image translation via latent energy transport

80：Unpaired image-to-image translation using adversarial consistency loss（ACLGAN）

相关工作：

GAN训练技术：

the Wasserstein loss的关键成分是一个梯度惩罚项，可以防止鉴别器梯度的大小变得太大。然而，Wasserstein loss后来被重新检验。值得注意的是，评估显示，稳定训练的是梯度惩罚项，而不是Wasserstien损失函数。此外，StyleGAN v2依赖于零中心梯度惩罚项，在高分辨率图像生成任务中实现最先进的结果。这些发现促使本工作探索应用梯度惩罚项来提高GAN训练稳定性。

计算机视觉的Transformer结构：

卷积神经网络(CNN)架构是计算机视觉任务的热门选择。在自然语言处理(NLP)领域，注意力机制和Transformer架构在开放基准任务中已经超越了以前的模型，如hidden Markov模型和循环神经网络。与CNN相比，Transformer可以更有效地捕获自然界中常见的非局部模式。而最近的其他工作[77, 35]表明CNN-Transformer混合可以获得更好的性能。

77: Early convolutions help transformers see better

35: Cmt: Convolutional neural networks meet vision transformers

自监督预训练（Self-supervised Pre-training）：

自我监督预训练通过在没有监督的情况下训练网络从原始数据中获得的人工任务来启动网络初始权重。这对于在小型标记数据集上训练具有大量参数的模型尤其重要，因为它们往往会过拟合。

我们假设GAN训练也可以从自我监督的预训练中受益。特别是GAN训练，已知会遭受“模式崩溃”问题:生成器无法准确重现图像的目标分布。观察发现，在开始GAN训练后的几个周期内，模式崩溃问题就发生了。这表明可以使用更好的初始化模型权重。事实上，GAN的迁移学习(transfer learning，一种预训练形式)已经成为提高GAN在小型训练数据集上性能的有效方法。我们只选择在自监督的修补任务上预训练生成器，这在NLP和计算机视觉中都是成功的。此外，它非常适合于图像到图像的转换模型，其中模型的输出形状与其输入形状相同。

方法：

CycleGAN类模型介绍

判别器通过与区分真实图像和翻译图像失败对应的反向传播损失(称为生成对抗损失或GAN损失)更新鉴别器:

生成器通过三个来源的反向传播损耗进行更新:GAN损失、周期一致性损失（cycle-consistency loss）和鉴别一致损失（identity-consistency loss）。以GA→B为例:

即：

为了改进原始的CycleGAN模型的性能，我们实现了三个主要更改。

首先，我们修改生成器，使其具有基于VIT和UNet的混合架构

其次，为了正则化CycleGAN鉴别器，我们增加了带有梯度惩罚项的vanilla CycleGAN判别器损失

最后，我们不再从随机初始化的网络权重中进行训练，而是在图像修补任务上以自监督的方式预训练生成器，以获得更好的启动状态。

UNet-ViT 生成器：

一个UNet-ViT生成器由一个UNet和瓶颈处的像素级Vison Transformer(ViT)组成：

UNet的编码路径通过四层卷积和下采样从输入中提取特征。每一层提取的特征也通过跳跃连接传递到解码路径的相应层，而最底部的特征则传递给VIT层。我们假设跳过连接是有效的传递高频特征到解码器，那么ViT为学习低频特征的成对关系提供了一种有效的方法。

在UNet的编码路径上，预处理层将图像转换为tensor(w0，h0，f0)，在每个下采样层之后，预处理tensor的高度和宽度减半，那么VIT层tensor的输入维度为(w, h, f) = (w0/16，h0/16，8f0)。

像素级的ViT(图2B)主要由Transformer编码器块的堆栈组成。为了构造堆栈的输入，ViT首先沿空间维度将编码图像压平，以形成token序列。token序列的长度为w × h，序列中的每个token都是一个长度为f的向量（即逐像素取token， 一共有w × h个token，每个token的长度为原tensor通道数f）。然后将每个token与维度为fp的二维傅里叶位置编码concat起来，并将结果线性映射为fv。为了提高Transformer的收敛性，我们采用rezero regularization的方案，并引入一个可训练的缩放参数α，调节残差块的重要分支的大小。Transformer堆栈的输出被线性投影回维为f，宽度为w和h。在本研究中，我们使用了12个Transform编码器块和并设置f，fp，fv 为384，且每个块中的前馈网络fh = 4fv。

梯度惩罚判别器损失(GP)：

在本研究中，我们在Eq.(1)-中使用最小二乘GAN (LSGAN)损失函数[50](即GAN损失是L2损失)，并用梯度惩罚项增强判别器损失。因此，我们决定采用更通用的梯度惩罚的方法，公式如下：

在实验中，这种以γ为中心的GP正则化提供了更稳定的训练，并且对超参数选择不那么敏感。

自监督预训练：

预训练是为下游任务准备大型网络的一种有效方法[24,7]，通常可以比随机初始化带来显著改善。在这项工作中，我们对UVCGAN生成器进行了图像修补任务的预训练。更准确地说，我们用大小为32 × 32的无重叠patch平铺（tile）图像，并通过将其像素值设置为零来遮盖40%的补丁。通过将像素值设置为0来mask掉40%的patches。生成器被训练为使用像素级L1损失预测原始未掩码图像。

UVCGAN训练步骤：

预训练：

为了构建受损图像，我们将大小为256 × 256的图像平铺成不重叠的32 × 32像素patches，并通过将像素值归零随机mask掉40%的patches。我们使用Adam优化器，余弦退火学习率衰减策略(cosine annealing learning-rate scheduler），以及一些标准数据增强，如小角度随机旋转，随机裁剪，随机翻转和色彩抖动(Color Jittering：对颜色的数据增强：图像亮度、饱和度、对比度变化)。在预训练过程中，我们不区分图像域，这意味着后续翻译训练中的两个生成器具有相同的初始化。

图像转换训练：

我们使用Adam优化器，在训练的前半段，学习率保持恒定在0.0001，然后在训练的后半段线性退火到零。我们应用了三种数据增强:调整大小、随机裁剪和随机水平翻转。

超参数搜索：

UVCGAN损失函数依赖于四个超参数: λcyc， λGP， λidt， γ，如果使用恒等损耗(λidt)，则总是按照CycleGAN原论文的建议设置为λcyc/2。为了找到性能最佳的配置，我们在网格上运行一个小规模的超参数优化。我们的实验表明，所有三个基准测试任务的最佳性能都是通过LSGAN + GP (λGP = 0.1， γ = 100)和在图像翻译数据集本身上预训练的生成器实现的。CelebA和Selfie2Anime的最佳λcyc分别在5和10。

评估指标：

Frechet Inception Distance (FID)和Kernel Incep-´tion Distance (KID)是两个最被接受的用于评估图像到图像转换性能的指标。分数越低，意味着翻译后的图像与目标域中的图像更相似。我们假设这是由于:

(1)随着数据样本越多，测试数据的大小不同, FID会降低，。

(2)测试前后处理的差异

(3)不同的度量公式(例如U-GAT-IT中的KID)

(4)不同的FID和KID实现。

因此，我们将评估标准化如下:

1)为FID和KID 使用完整测试数据集，对于KID子集Selfie2Anime使用50，两个CelebA数据集使用1000;

2)调整非方形CelebA图像的大小，并采用大小为256 × 256的中央裁剪来保持正确的纵横比;

3)将所有KID和FID计算用torch-fidelity package。

（代码链接：GitHub - toshas/torch-fidelity: High-fidelity performance metrics for generative models in PyTorchHigh-fidelity performance metrics for generative models in PyTorch - GitHub - toshas/torch-fidelity: High-fidelity performance metrics for generative models in PyTorchhttps://github.com/toshas/torch-fidelity）

注意力模块的说明

注意矩阵的(i, j)项表明当第i行的和为1时，token i放了多少注意力在token j上。当使用多头注意力时，每个头产生一个注意矩阵。为了简单起见，我们平均Transformer编码器堆栈中每个块的所有头部和目标token的注意力权重。给定大小为256 × 256的输入图像，这提供了一个维数为w × h (16 × 16 = 256)的注意向量。这样一个向量的第j个项表示token j平均收到多少关注。由于所有tokens代表原始图像的不重叠patches，我们将用以下步骤生成注意力heatmap：reshape所有特征向量为16×16的方形，将其放大16倍以匹配输入图像的尺寸，然后应用σ = 16的高斯滤波器。通过将注意力热图叠加在输入图像上，我们注意到每个块都在关注特定的面部部位，其中眼睛和嘴巴区域受到的关注最多。

附录：

通过消融实验发现，没有GP项的自监督生成器预训练似乎并没有提高图像到图像的翻译性能，有时还会使其更差，自监督的预训练只有在与GP一起使用时才有帮助。

UNet-ViT生成器的细节：

ViT结构细节：