摘要
(1)条件生成对抗网络(cGAN)为许多计算机视觉和图形应用程序提供了可控制的图像合成。 但cGAN的计算强度要比用于识别的CNN高1-2个数量级。
- 例如,Gau GAN每个图像消耗281G MACs,而MobileNet-v3只消耗0.44G MACs,这使得交互式部署变得非常困难。
(2)在本文中,作者提出了一个通用的压缩框架来减少cGANs中生成器的inference time和model size。直接应用现有的CNNs压缩方法会产生较差的性能,这是由于GAN训练的难度和生成器架构的差异造成的。
(3)作者以两种方式应对这些挑战:
- 首先,为了稳定GAN训练,我们将原模型的多个中间表示的知识转移到其压缩模型中,将非成对和成对学习统一起来。
- 其次,我们的方法不是重用现有的CNN设计,而是通过神经架构搜索(neural architecture search, NAS)自动找到有效的架构。为了加速搜索过程,我们通过权值共享来解耦模型训练和架构搜索。
(4)实验证明了我们的方法在不同的监督设置(配对和不配对)、模型架构和学习方法(如pix2pix、GauGAN、CycleGAN)下的有效性。在不损失图像质量的前提下,我们将CycleGAN的计算量减少了20×以上,GauGAN的计算量减少了9×。
1. Introduction
生成式对抗网络(GANs)擅长合成逼真的照片。conditional GANs允许可控的图像合成。然而,边缘设备(移动电话、平板电脑、VR耳机)受到内存和电池等硬件资源的严格限制。这种计算瓶颈阻碍了conditional GANs在边缘设备上的部署。
- 与图像识别用到的CNNs不同的是,基于图像的conditional GANs的计算量非常大。例如,CycleGAN模型需要超过50G MACs,比MobileNet多100倍。近期的模型GauGAN,虽然能生成逼真的高分辨率图像,但需要250G MACs。
为了解决这些问题:
- 首先:将知识从原始teacher生成器的中间表示层迁移到其相应的压缩student生成器层。作者还发现,使用teacher模型的输出为非配对训练创建伪配对是有益的,这将把非配对学习转化为配对学习。
- 其次,我们使用神经架构搜索(neural architecture search, NAS)来自动找到一个计算成本和参数显著减少的高效网络。为了降低训练成本,我们通过训练一个包含所有可能的channel数量配置的“once-for-all”网络,将模型训练从架构搜索中解耦出来。once-for-all 网络可以通过权值共享产生多个子网络,不需要再训练就可以评估每个子网络的性能。作者的方法可以应用于各种cGAN模型,无论考虑模型架构、学习算法和监督设置(配对或不配对)。
通过大量的实验,作者的方法可以减少三种广泛使用的条件GAN模型的计算量,包括pix2pix, CycleGAN , GauGAN,在MACs上减少9× ~ 21×的计算量,同时不丢失生成图像的视觉保真度,如图1所示:
图1
2. Related Work
2.1 Conditional GANs.
生成式对抗网络(GANs)擅长合成逼真的照片结果。它的条件形式conditonal GANs进一步支持可控的图像合成,允许用户合成给定各种条件的输入的输出图像。后续的工作进一步提高了结果的分辨率和真实感。后来,提出了几种unpaired data学习cGANs的算法。
图像超分辨、真实感图像的合成结果需要大量的计算。如图2所示:
图2
- 尽管模型大小与图像识别CNNs的大小相同,但cGANs需要比其它大两个数量级的计算量。这使得在有限的计算资源下将这些模型部署到边缘设备非常具有挑战性。本文作者关注用于交互式应用( interactive applications)的高效的 image-conditional GANs架构。
2.2 Model acceleration.
- 在各种实际应用中,硬件高效的深度学习受到了广泛的关注。为了减少网络权重的冗余,研究人员提出了修剪各层之间的连接。然而,经过修剪的网络需要专门的硬件来实现它的完全加速。
- 随后的几项工作提出了修剪整个卷积滤波器(convolution filter),以提高计算的规律性。AutoML for Model Compression (AMC)利用强化学习来自动确定每一层的修剪比例。Liu等人后来用进化搜索算法代替了强化学习。最近,Shu等人通过对原有的CycleGAN算法进行了修改,提出了CycleGAN的协同进化剪枝(co-evolutionarypruning)方法。这种方法是为特定的算法量身定做的。在适当的压缩比(4.2×)下,压缩模型显著提高了FID。
- 相比之下,本文的模型无关方法可以应用于具有不同学习算法、体系结构以及成对和非成对设置的cGANs。本文假设不知道原始的cGAN学习算法。实验表明,我们的通用方法在保留原模型FID的同时,达到了21.1×压缩比(相比于比特定于CycleGAN的方法好5倍)。
2.3 Knowledge distillation.
Hinton等人提出了知识蒸馏(Knowledge distillation),将更大的teacher网络中的知识转移到更小的student网络中,student网络被训练成模仿teacher网络的行为。
- 有几种方法利用知识蒸馏来压缩识别模型。最近,Aguinaldo等人采用了这种方法来加速无条件的GANs。与它们不同的是,我们关注的是cGANs。我们在cGANs上实验了几种蒸馏方法,但只观察到有限的改进,不足以用于交互式应用。
2.4 Neural architecture search.
神经结构搜索(Neural Architecture Search, NAS)成功地设计出了在大规模图像分类任务中表现优于手工制作的神经网络结构。为了有效降低搜索成本,研究人员最近提出了一种once-for-all神经结构搜索,其中不同的候选子网络可以共享同一组权值。虽然所有这些方法都侧重于图像分类模型,但作者使用NAS研究了高效的cGANs体系结构。
3. Method
压缩交互式应用程序的条件生成模型具有挑战性,原因有二:
- 首先,GANs的训练动态本质上是高度不稳定的。
- 其次,识别模型和生成模型在架构上的巨大差异使得现有的CNN压缩算法难以直接应用。
针对上述问题,作者提出了一种适合高效生成模型的训练协议(第3.1节),并通过神经结构搜索(NAS)进一步提高压缩比(第3.2节)。整个框架如图3所示:
图3
3.1. Training Objective
3.1.1 Unifying unpaired and paired learning.
cGANs的目标是学习源域X和目标域Y之间的映射函数G:
- 它们使用配对数据 { x i , y i } i = 1 N \{x_i, y_i\}_{i=1}^N {xi,yi}i=1N,其中 x i ∈ X x_i∈X xi∈X 和 y i ∈ Y y_i∈Y yi∈Y。 或 使用未配对数据(源数据集 { x i } i = 1 N \{x_i\}_{i=1}^N {xi}i=1N到目标数据集 { y j } j = 1 M \{y_j\}_{j=1}^M {yj}j=1M)进行训练。这里,N和M表示训练图像的数量。为了简单起见,我们省略了下标i和j。
已经提出了几个学习目标来处理paired和unpaired settings。广泛的训练目标使得构建通用的压缩框架非常困难。为了解决这个问题,作者在模型压缩设置中统一了非成对和成对的学习,而不管teacher模型最初是如何训练的。
- 给定原始的teacher 生成器 G ′ G' G′,我们可以将unpaired的训练集转化为paired的训练集。特别地,对于unpaired的设置,我们可以将原始生成器的输出视为ground-truth,并使用成对的学习目标来训练我们的压缩生成器G。我们的学习目标可以总结如下:
通过这些修改,我们可以将相同的压缩框架应用于不同类型的cGANs。此外使用上述伪对进行学习,与原始的非成对训练设置相比,训练更加稳定,效果也更好,此时未配对训练已经转化为配对训练。
3.1.2 Inheriting the teacher discriminator.
虽然我们的目标是压缩生成器generator,但判别器 D D D在识别当前生成器的弱点时,会存储已学习GAN的有用知识。
- 因此,作者采用相同的判别器结构,使用teacher预先训练的权重,并与我们的压缩的生成器一起对判别器进行fine-tune。在实验中,我们观察到一个预先训练的判别器可以指导我们的student 生成器的训练。使用一个随机初始化的判别器往往会导致严重的训练不稳定性和图像质量的下降。GAN的目标函数为:
在这里,我们使用来自teacher判别器 D ′ D' D′的权重来初始化student判别器 D D D。G和D使用标准的极小极大优化进行训练。
3.1.3 Intermediate feature distillation.
知识蒸馏是一种广泛使用的CNN模型压缩方法。通过匹配输出层的对数分布,可以将 dark knowledge从 teacher 模型迁移到 student 模型,提高 student 的性能。
- 然而,cGANs通常输出的是确定性的图像,而不是概率分布。因此,很难从教师的输出像素中提取出 dark knowledge。
- 特别是在成对训练的情况下,teacher 模型生成的输出图像与真实目标图像相比,基本不包含额外的信息。实验表明,在配对训练中,单纯地模仿教师模型的输出并没有带来改进。
为了解决上述问题,作者匹配 teacher 生成器的中间表示,像之前的工作所探索的那样,中间层包含更多的channels,提供更丰富的信息,并允许 student 模型获得除输出之外的更多信息。蒸馏目标函数可以形式化为:
其中, G t ( x ) G_t(x) Gt(x)和 G t ′ ( x ) G'_t(x) Gt′(x)是 student 和 teacher 模型中第T个选择层的中间特征激活,T表示层数。1×1可学习卷积层 f t f_t ft将 student 模型的特征映射到 teacher 模型特征中相同数量的channels。我们共同优化 G t G_t Gt和 f t f_t ft,使蒸馏损失最小化。附录6.1详细说明了我们在实践中选择的层。
3.1.4 Full objective.
我们的最终目标如下:
其中超参数 λ r e c o n λ_{recon} λrecon和 λ d i s t i l l λ_{distill} λdistill控制每个项的重要性。
3.2. Efficient Generator Design Space
选择一个设计良好的 student 架构对于知识蒸馏的最终表现是至关重要的。我们发现,单纯地缩小 teacher 模型的通道数并不能产生一个紧凑的 student 模型:当计算减少超过4倍时, 性能开始显著下降。其中一个可能的原因是现有的生成器架构通常采用图像识别模型,可能不是图像合成任务的最佳选择。
下面从现有的cGAN生成器派生出更好的架构设计空间,并在该空间内执行神经架构搜索(NAS)。
3.2.1 Convolution decomposition and layer sensitivity.
现有的生成器通常采用vanilla卷积来遵循分类和分割CNNs的设计。最近高效的CNN设计广泛采用了卷积的分解版本(depthwise+pointwise),这证明了有更好的性能计算权衡。我们发现使用分解的卷积也有利于cGANs中生成器的设计。
然而,通过一些实验已经表明,直接对所有卷积层进行分解(如在分类器中)将显著降低图像质量。分解一些层会立即影响性能,而其他层则更健壮。此外,这个层敏感性模式也不同于识别模型。
- 例如,在ResNet生成器中,resBlock层消耗了大部分的模型参数和计算成本,而几乎不受分解的影响。相反,上采样层有更少的参数,但是对模型压缩相当敏感:中等程度的压缩会导致很大的FID退化。因此,我们只分解resBlock层。我们在第4.4节中对层的敏感性进行了全面的研究。
3.2.2 Automated channel reduction with NAS.
现有的生成器在所有层上使用手工制作的(而且基本上是统一的)通道数,其包含冗余信息,远远不是最优的选择。为了进一步提高压缩比,我们使用通道剪枝(channel pruning)方法来自动选择生成器中的通道宽度来去除冗余,这样可以二次减少计算量。我们支持关于通道数量的细粒度选择。对于每个卷积层,通道的数量可以从8的倍数中选择,这平衡了MACs和硬件并行性。
给定可能的通道配置 c 1 , c 2 , … , c K {c_1,c_2,…,c_K} c1,c2,…,cK,其中K是要修剪的层数,我们的目标是使用神经结构搜索(NAS)找到最佳的通道配置:
其中 F t F_t Ft为计算约束。
- 一种直接的方法是遍历所有可能的通道配置,训练它至收敛,然后评估并选择性能最好的生成器。
- 然而,随着K的增加,可能的配置的数量呈指数级增长,每个配置可能需要不同的超参数,这些超参数与每个项的学习率和权重有关。这种尝试和错误的过程太耗时了。
3.3. Decouple Training and Search
为了解决这个问题,作者将模型训练从架构搜索中解耦出来,就像最近once-for-all神经架构搜索方法的工作一样。
- 我们首先训练一个支持不同channels数量的 once-for-all 网络。每个通道数目不同的子网络都经过同等的训练,可以独立运行。子网络与“once-for-all”网络共享权重。图3说明了整个框架。我们假设原始的 teacher 生成器有 { c k 0 } k = 1 K \{c_k^0\}_{k=1}^K {ck0}k=1K个通道。对于一个给定的通道数量配置 { c k } k = 1 K \{c_k\}_{k=1}^K {ck}k=1K , c k ≤ c k 0 c_k≤c_k^0 ck≤ck0,我们按照Guo等人的方法,从“once-for-all”网络的相应权张量中提取第一个 { c k } k = 1 K \{c_k\}_{k=1}^K {ck}k=1K的通道,得到子网络的权值。
- 在每个训练步骤,我们随机采样带有某通道数量配置的子网络,计算出输出和梯度,使用我们的学习目标函数(等式4)提取和更新权重。由于权重在最初的几个通道更新更加频繁,他们在所有的权重中发挥更重要的作用。
Once-for all 论文翻译
Once-for all 论文解读
- 在对“once-for-all”网络进行训练后,我们通过直接评估验证集上每个候选子网络的性能来找到最佳的子网络。由于“once-for-all”网络经过了充分的训练,并进行了权值共享,因此不需要进行微调。
- 当从零开始训练时,这近似于模型的性能。通过这种方式,我们可以解耦生成器架构的训练和搜索:我们只需要训练一次,但是我们可以评估所有可能的通道配置,而不需要进一步的训练,并选择最好的一个作为搜索结果。我们还可以选择微调所选的体系结构,以进一步提高性能。
4. Experiments
4.1. Setups
4.1.1 Models.
我们对三个cGAN模型进行了实验:
- CycleGAN:用unpaired数据实现图像到图像之间的翻译模型,使用基于ResNet的生成器。
- Pix2Pix:基于cGAN的paierd图像到图像翻译模型,使用了基于U-Net的生成器。 对于该模型,我们观察到在给定相同学习目标的情况下,基于ResNet的生成器以更少的计算成本实现了更好的结果,因此,我们将基于ResNet的生成器替换了原始的U-Net生成器。
- GauGAN:是最先进的paired图像翻译模型。 给定语义标签图,它可以生成高保真图像。
作者使用经过上述修改的官方PyTorch存储库重新训练了pix2pix和CycleGAN。
- 经过重新训练的pix2pix和CycleGAN模型(在我们的仓库中提供)略胜于官方预训练的模型。 我们将这些经过重新训练的模型用作原始模型。
- 对于Gau GAN,我们使用作者的预训练模型。
4.1.2 Datasets.
我们使用以下四个数据集:
- Edges→shoes: 我们使用来自UT Zap pos50K数据集的69,025张图像。 我们随机分割数据集,以便验证集具有2,048张图像,以便稳定地评估Frechet起始距离(FID)(请参阅第4.2节)。 我们在此数据集上评估pix2pix模型。
- Cityspaces: 数据集包含德国人街道场景的图像。 训练集和验证集分别包含2975和500张图像。 我们在此数据集上评估pix2pix和GauGAN模型。
- Horse↔zebra:该数据集由最初来自ImageNet [13]的1187马图像和1,474斑马图像组成,并在CycleGAN [76]中使用。 验证集包含120个马图像和140个斑马图像。 我们在该数据集上评估CycleGAN模型。
- Map↔aerial photo(航拍照片):该数据集包含从Google Maps抓取得到,并用于pix2pix的2194张图像。训练集和验证集包含1096和1098张图片。 我们在此数据集上评估pix2pix模型。
4.1.3 Implementation details.
- 对于CycleGAN和pix2pix模型,在所有实验的训练过程中,我们对生成器和判别器均使用0.0002的学习率。
- 数据集horse→zebra, edges→shoes,map→aerial photo, and cityscapes的batch_size大小分别为1、4、1和1。 对于GauGAN模型,我们遵循原始论文中的设置,不同之处在于批量大小为16而不是32。我们发现以较小的批量大小可以获得更好的结果。
4.2. Evaluation Metrics
介绍用于评估生成图像相等性的指标:
Frechet Inception Distance (FID):
FID分数旨在计算使用InceptionV3 网络从真实图像和生成图像中提取的特征向量的分布之间的距离。 该分数用于衡量真实图像与生成图像之间的相似度。 较低的分数表示生成的图像质量较好。 我们使用开源的FID评估代码。
- 对于paired的图像到图像转换(pix2pix和Gau GAN),我们计算了转换后的测试图像与实际测试图像之间的FID。
- 对于unpaired的图像到图像翻译(CycleGAN),我们计算翻译后的测试图像与实际训练+测试图像之间的FID。 就像以前的无条件GAN一样,这使我们可以使用更多图像进行稳定的FID评估。 当我们使用真实的测试图像而不是真实的训练+测试图像时,压缩的CycleGAN模型的FID会稍微增加。
Semantic Segmentation Metrics:
我们采用语义分割度量来评估Cityscapes数据集上生成的图像。 我们在生成的图像上运行语义分割模型,并比较分割模型的执行效果。 我们选择(mean Intersection overUnion, mIoU)作为细分指标,并使用DRN-D-105作为细分模型。 更高的信号强度意味着所生成的图像看起来更逼真并且可以更好地反映输入的标签图。 我们将DRN-D-105的输出语义图上采样为2048×1024,这是Cityscapes地面真实图像的分辨率。
DRN-D-105(Dilated Residual Networks)