当前位置: 首页 > 科技观察

韩松、朱俊彦等人提出了GAN压缩方式:算力消耗不到1-9,开源的

时间:2023-03-22 16:11:35 科技观察

GAN也可以进行大幅压缩。麻省理工学院韩松团队的最新研究启发了许多研究人员。生成模型GAN是机器学习领域最重要的发展方向之一。然而,这类算法消耗的计算能力巨大,大多数研究人员很难得出新的结果。近年来,这个方向有被大机构垄断的趋势。但最近,来自麻省理工学院(MIT)、Adobe和上海交通大学的研究人员提出了一种压缩条件GAN的通用方法。这种新技术将广泛使用的条件GAN模型如pix2pix、CycleGAN和GauGAN的计算量减少到1/9~1/21,同时保持视觉保真度。该方法适用于各种生成器架构、学习目标、配对或非配对设置。目前该研究论文已被CVPR2020会议收录,GAN压缩框架的PyTorch实现也已开源。项目链接:https://github.com/mit-han-lab/gan-compressionGAN压缩后具体表现如何?在研究人员展示的演示中,使用CycleGAN为视频中的马添加斑马条纹需要不到1/16的计算能力,帧数增加了三倍,效果也得到了提升:它值得一提的是,研究所使用的硬件平台是英伟达的边缘AI计算芯片JetsonXavierGPU。官方数据显示,JetsonXavier的INT8算力为22+10TOPS,骁龙865为15TOPS。压缩后的GAN现在似乎可以运行在机器人、无人机等小型设备上,未来会塞进手机里。论文链接:https://arxiv.org/pdf/2003.08936v1.pdf一般介绍生产对抗网络(GANS)擅长合成非常逼真的图像。GAN的一种变体,条件生成对抗网络(cGAN),可以在许多计算机视觉和成像应用程序中实现可控的图像合成。然而,这些应用中的大多数都需要模型来与人交互,因此需要低延迟设备来获得更好的用户体验。然而,最近引入的一些cGAN的计算强度比当前识别卷积神经网络(CNN)高1到2个数量级。例如,GanGAN每张图片消耗281GMac,而MobileNet-v3只需要0.44GMac,这使得前者难以用于交互式部署。而且,就目前的情况来看,大多数边缘设备都受到内容和电池等硬件的限制,这也阻碍了GAN在边缘设备上的部署。因此,基于GAN和cGAN在图像合成领域存在的这些问题,韩松团队提出了GAN压缩,这是一种通用的压缩方法,旨在减少GAN的推理时间和计算成本。同时,压缩生成模型面临两个基本困难:GAN训练不稳定,尤其是未配对时;生成器与CNN不同,因此很难使用现有的CNN设计。为了解决这个问题,该团队将知识从原来的教师生成器中间表示层转移到其对应的学生生成器层。为了降低训练成本,该团队还通过训练一个包含所有可能通道数的“once-for-all网络”,将模型训练与架构搜索分开。这个“once-for-allnetwork”可以通过权值共享生成很多子网络,每个子网络的性能都可以在不训练的情况下进行评估。该模型可以应用于各种条件下的GAN模型,无论模型架构、学习算法或监督设置(配对或非配对)如何。通过大量实验,该团队证明该方法可以将三种广泛使用的GAN模型pix2pix、CycleGAN和GauGAN的计算量减少到1/9到1/21,而不会损失生成图像的保真度。众所周知,压缩交互式应用程序的条件生成模型具有挑战性,主要原因有两个。首先,GAN的动态训练本质上是非常不稳定的;其次,识别模型和生成模型之间的巨大架构差异使得难以直接使用现有的CNN压缩算法。基于这些原因,研究人员提出了专门针对高效生成模型的训练方案,并使用神经架构搜索(NAS)进一步提高压缩率。GAN压缩框架的整体架构如下图3所示,他们以ResNet生成器为例。应该强调的是,相同的框架可以应用于不同的生成器架构和学习目标。图3:本文GAN压缩框架的整体架构图。ObjectiveFunction1.UnifiedPairedandUnpairedLearning广泛的训练目标使得难以构建通用的压缩框架。为了解决这个问题,我们在模型压缩设置中统一了配对和非配对学习,而不管教师模型最初是如何训练的。给定原始教师生成器G',研究人员将未配对的训练设置转换为配对设置。对于未配对的设置,可以将原始生成器输出视为真实值,并使用配对训练目标训练压缩生成器G。学习目标总结如下:基于这些修改,现在可以将相同的压缩框架应用于不同类型的cGAN。此外,使用上述伪对进行学习会导致更稳定的训练,并产生比原始未配对训练设置更好的结果。2.从教师判别器中学习虽然这项研究的重点是压缩生成器,但对GAN有用的信息存储在判别器D中。因此,我们利用相同的判别器架构,使用预训练权重对判别器和压缩生成器一起进行微调来自教师鉴别器。在实验中,研究人员观察到预训练的判别器可以指导学生生成器的训练。使用随机初始化的鉴别器通常会导致训练不稳定和图像质量差。这个GAN的目标可以写成如下:在上面的等式中,研究人员用来自教师鉴别器D'的权重初始化学生鉴别器D。他们使用标准的极小极大优化器训练G和D。3.中间特征蒸馏知识蒸馏是CNN模型压缩中常用的方法。通过匹配输出层logit的分布,可以将teacher模型的暗知识迁移到student模型,提高student模型的性能。然而,条件GAN通常输出确定性图像而不是概率分布。为了解决上述问题,研究人员匹配了教师生成器的中间表示。中间层包含的通道越多,它们提供的信息就越多,学生模型在输出之外可以捕获的信息就越多。蒸馏目标如下:其中,G_t(x)和G′_t(x)是学生和教师模型中第t个选定层的中间特征激活,T表示层数。4.完整的优化目标最终目标可以写成如下:其中超参数λ_recon和λ_distill控制每一项的重要性。高效的生成器设计空间选择一个设计良好的学生架构对于最终的知识蒸馏效果至关重要。研究人员发现,简单地减少教师模型的通道数并不能使学生模型更加紧凑:当计算量减少超过4时,性能会显着下降。1.卷积分解和层敏感性现有的生成器通常采用传统的卷积来匹配CNN分类和分割设计。最近,一些高效的CNN设计广泛采用了卷积的分解形式(depthwise+pointwise),在性能和计算上有更好的平衡。研究人员发现,分解卷积也可以用于cGAN的生成器设计。2.使用NAS自动裁剪通道现有的生成器在所有层中使用手动设计(并且几乎统一)的通道数量,这会产生冗余并且远非最佳。为了进一步提高压缩效率,研究人员使用通道剪枝在生成器中自动选择通道宽度,从而减少冗余,减少两倍的计算量。这种方法支持通道数的细粒度选择,并且对于每个卷积层,可以从8的倍数中选择卷积层,平衡MAC和硬件并行性。解耦训练和架构搜索研究人员遵循最近关于一次性NAS方法的工作,将模型训练与架构搜索解耦。首先,训练一个支持不同数量通道的“一次性”网络,其中每个子网络均等训练。图3说明了总体框架。研究人员假设原始教师生成器有一个通道,对于给定的通道数,从“once-for-all”权重张量中提取第一个通道,得到子网络的权重网络。在每个训练步骤中,使用学习目标随机采样具有一定数量通道的子网络,计算输出和梯度,并更新提取的权重(等式4)。由于提取的前几个通道更新更频繁,因此它们在所有权重中起着更关键的作用。在训练完这个“once-for-all”网络后,研究人员直接在验证集上评估每个子网络的性能,以找到最好的子网络。“once-for-all”网络已经通过权重共享进行了彻底的训练,不需要微调。该结果近似于从头开始训练的模型的性能。这样,只需要训练一次,无需进一步训练,就可以评估所有通道的所有配置,并根据搜索结果找到最好的配置。当然,还可以对所选架构进行微调,进一步提升其性能。实验结果研究人员对以下三种条件GAN模型进行了实验,以验证本文GAN压缩框架的泛化性,它们是CycleGAN、Pix2Pix和GauGAN。使用的四个数据集是Horse?zebra、Edges→shoes、Cityscapes和Map?aerialphoto。下面的表1给出了在上述四个数据集上压缩CycleGAN、Pix2Pix和GauGAN模型的量化结果。表1:在Cityscapes数据集上使用mAP度量(越高越好)和在其他数据集上使用FID度量对三种条件GAN模型压缩进行定量评估。结果表明,我们的GAN压缩方法可以将当前SOTA条件GAN在MAC上压缩7到21倍,在模型大小上压缩5到33倍,性能略有下降。对于CycleGAN模型压缩,本文的GAN压缩方法远优于之前的CycleGAN-specificCo-evolution方法。性能和计算之间的权衡这种方法除了可以实现较大的压缩比外,还可以提高不同模型尺寸的性能。下面的图6显示了pix2pix模型中不同数据集的性能与计算权衡。图6:pix2pix在Cityscapes和Edges→Shoes数据集上的权衡曲线。修剪和蒸馏方法在大型模型中优于从头开始的训练,但在模型被急剧压缩时表现不佳。效果展示下面图4展示了使用该方法得到的效果。图中分别显示了输入数据、标准输出、原始模型输出和压缩模型输出。从图中可以看出,即使在较大的压缩比下,研究人员提出的方法仍能保持输出图像的视觉可信度。图4:Cityscapes、Edges→Shoes和Horse→Zebra数据集的性能比较。硬件推理加速对于实景交互应用,推理加速对硬件设备的重要性远大于降低计算消耗。如下表2所示,为了验证所提方法在实际应用中的有效性,研究人员在不同计算性能的设备上测试了压缩模型的推理速度。表2:在NVIDIAJetsonAGXXavier、NVIDIAJetsonNano、1080TiGPU和XeonCPU上测得的内存和延迟下降。结论在本文中,韩松团队提出的通用压缩框架可以显着降低条件GAN中生成器的计算成本和模型大小,并通过知识蒸馏和自然架构搜索提高训练稳定性和模型效率。实验表明,所提出的GAN压缩方法能够在保持视觉质量的同时压缩多个条件GAN模型。他们表示,未来的研究工作将集中在减少模型延迟和为生成视频模型构建高效框架上。