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

发布 : 2021-02-14 分类 : 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的图像分辨率训练时。
image.png
在本文中,我们的目标是在高分辨率图像上学习无条件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模型小得多,而且训练起来也快得多。同时,由于采用两种所提出技术的紧凑型模型,我们的模型在小型数据集上仍然保持稳健。
image.png
image.png

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的优点,并且没有额外的计算负担。
image.png
尽管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)\|]
$$
image.png
我们自监督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}
$$

学习、记录、分享、获得

微信扫一扫, 向我投食

微信扫一扫, 向我投食