致力于更快更稳定的GAN训练 以实现高精确度的少样本图像合成

The datasets and code are available at: https://github.com/odegeasslbc/FastGAN-pytorch

摘要

在高保真图像上训练生成对抗网络（GAN）通常需要大规模的 GPU 集群和大量的训练图像。在本文中，我们以最低的计算成本研究了 GAN 的少量图像合成任务。我们提出了一种轻巧的 GAN 结构，该结构在 1024×1024 分辨率下可获得卓越的质量。值得注意的是，该模型仅需在单个 RTX-2080 GPU 上进行几个小时的培训即可从头开始收敛，即使在少于 100 个训练样本的情况下，也具有一致的性能。两种技术设计构成了我们的工作，一个是跳过层的通道激励模块，另一个是经过自我监督的，被识别为特征编码器的鉴别器。在数据和计算预算有限的情况下，在 13 个数据集，涵盖了广泛的图像域，与最先进的 StyleGAN2 相比，我们展示了该模型的优越性能。

引言

在现实生活中，用于训练 GAN 的可用样本可能很少，例如罕见疾病的医学图像，特定名人的肖像集和特定艺术家的艺术品。使用预先训练的模型进行转移学习（Mo 等人，2020； Wang 等人，2020）是缺乏训练图像的一种解决方案。然而，不能保证找到一个兼容的预训练数据集，此外，如果没有，微调可能会导致更差的性能（Zhao et al。，2020）。
在最近的一项研究中，强调了在艺术创作应用程序中，大多数艺术家更喜欢根据自己的图像从头开始训练模型，以避免微调的预训练模型产生偏差。此外，研究表明，在大多数情况下，艺术家都希望使用少于 100 张图像的数据集来训练他们的模型（Elgammal 等，2020）。动态数据增强（Karras 等，2020a; Zhao 等，2020）用更少的图像平滑了差距并稳定了 GAN 训练。然而，SOTA 模型（例如 StyleGAN2（Karras 等人，2020b）和 BigGAN（Brock 等人，2018））的计算成本仍然很高，尤其是在以 1024×1024 的图像分辨率训练时。

在本文中，我们的目标是在高分辨率图像上学习无条件 GAN，计算成本低，训练样本少。如图 2 所示，这些训练条件使模型面临过度拟合和模式崩溃的高风险 (Arjovsky & Bottou, 2017; Zhang & Khoreva, 2018)。在给定苛刻的训练条件下训练 GAN，我们需要能够快速学习的生成器 G 和可以不断提供有用信号来训练 G 的鉴别器 D。为解决这些挑战，我们总结出以下贡献：

• 我们设计了 Skip-Layer channel-wise Excitation（SLE）模块，该模块利用低比例激活来修改高比例特征图上的通道响应。SLE 可在整个模型权重中提供更鲁棒的梯度流，以加快训练速度。它还可以自动学习样式/内容分解，例如 StyleGAN2。
• 我们提出了一种自我监督的鉴别器 D，该鉴别器 D 被训练为具有额外解码器的特征编码器。我们迫使 D 从输入图像中学习更多描述性的特征图，覆盖更多区域，从而产生更全面的信号来训练 G。我们测试了 D 的多种自我监督策略，其中我们证明自动编码效果最好。
• 我们基于两种提出的技术构建了一个计算效率高的 GAN 模型，并在多个高保真数据集上展示了该模型的鲁棒性，如图 1 所示。

方法

我们为模型采用了简约的设计。特别是，我们在 G 的每个分辨率上使用单个 conv 层，并且在 G 和 D 的高分辨率（≥512×512）的 conv 层上仅应用三个（输入和输出）通道。图 3 和图 4 说明了我们的 G 和 D 的模型结构，并描述了组件层和前向流。这些结构设计使我们的 GAN 比 SOTA 模型小得多，而且训练起来也快得多。同时，由于采用两种所提出技术的紧凑型模型，我们的模型在小型数据集上仍然保持稳健。

SKIP-LAYER CHANNEL-WISE EXCITATION

作者认为，为了合成更高分辨率的图像，生成器 G 不可避免的需要被加深，更深的模型需要更长时间去训练，由于模型参数的大量增加，通过 G 的梯度流也会被弱化。为了更好的训练模型，残差结构 ResBlock 被提出用于增强层之间的梯度信号。ResBlock 被广泛使用在 GAN 的结构，但是它也增加了计算成本。
利用跳跃连接的思想，我们将两种独特的思想重新设计出 Skip-Layer Excitation module (SLE)。首先，ResBlock 实现的跳跃连接需要不同卷积层的激活之间进行 element-wise 的加法，它需要空间维度相同。取而代之的是在不同激活之间采用 channel-wise 的乘法，从而消除了繁琐的卷积计算。其次，在之前的工作中，跳跃连接只能在相同分辨率内使用。相比之下，由于不在需要相同空间维度的限制，我们可以在分辨率更大的范围之间进行跳跃连接。
这两个设计使得 SLE 继承了 ResBlock 的 shortcut gradient flow 的优点，并且没有额外的计算负担。

尽管 SLE 像 SE 一样带来了通道方式的特征重新校准的好处，但它也像 ResBlock 一样增强了整个模型的梯度流。
通过从另一个合成样本中替换 SLE 中的$X_{low}$，我们的 G 可以生成一个内容不变的图像，但样式与新替换图像相同。

SELF-SUPERVISED DISCRIMINATOR

我们将 D 视为编码器，并使用小型解码器对其进行训练。这种自动编码训练迫使 D 提取图像特征，解码器可以给出良好的重构。解码器与 D 一起在简单的重建 loss 上进行了优化，仅在真实样本上进行训练：

$$ \mathcal{L}_{\text {recons }}=\mathbb{E}_{\mathbf{f} \sim D_{\text {encode }}(x), x \sim I_{\text {real }}}[\|\mathcal{G}(\mathbf{f})-\mathcal{T}(x)\|] $$

我们自监督 D 如上图所示，在这里使用两个解码器来处理两个尺度的特征图，$f_{1}$ on $16^{2}$ and $f_{2}$ on $8^{2}$。在$128x128$分辨率，解码器只有 4 个 卷积层，几乎不会引起额外的计算（比其他正则化方法少得多）。我们$\frac{1}{8}$高度和宽度，随机裁剪$f_1$，然后在相同位置裁剪 real image 对应的区域得到$I_{part}$。调整 real image 大小得到$I$。解码器从裁剪后的$f_1$得到$I^{'}_{part}$，从$f_2$得到$I^{'}$。 最后，将 D 和解码器一起训练，以最大程度降低式 2 的 loss。

这样的重构训练确保 D 从输入中提取出更全面的表示，涵盖总体构图（来自 f2）和详细纹理（来自 f1）。注意 G 和 T 不被限制为裁剪，更多高效操作仍待探索。auto-encoding 方法是被自监督学习采用的典型方法，它可以提高模型的鲁棒性和泛化性能。在 GAN 的背景下，我们发现通过自我监督训练策略进行规则化的 D 可以显着提高 G 的合成质量，其中自动编码带来了最大的性能提升。
尽管我们针对 D 的自我监督策略采用自动编码器（AE）的形式，但是这种方法与尝试将 GAN 和 AE 结合在一起的工作从根本上不同。后者的工作主要是在从 D 获得的潜在空间上训练 G 作为编解码器，或者将 AE 的对抗训练视为 D 的对抗性损失。相反，我们的模型是具有简单得多的训练模式的纯 GAN。自动编码训练仅用于正则化 D，其中不涉及 G。
总结来说，我们采用对抗损失的 hinge version（Lim & Ye (2017)），去迭代训练 D 和 G。我们发现，不同的 GAN loss 对性能的影响不大，而 hinge loss 的计算速度最快：

$$ \begin{aligned} &\mathcal{L}_{D}=-\mathbb{E}_{x \sim I_{\text {real }}}[\min (0,-1+D(x))]-\mathbb{E}_{\hat{x} \sim G(z)}\left[\min (0,-1-D(\hat{x})]+\mathcal{L}_{\text {recons }}\right.\\ &\mathcal{L}_{G}=-\mathbb{E}_{z \sim \mathcal{N}}[D(G(z))] \end{aligned} $$