生成模型 GAN 是机器学习领域里最为重要的发展方向之一。但这类算法需要消耗巨量算力,大多数研究者已经很难得出新成果。近年来,这一方向颇有被大型机构垄断的趋势。 但近日,来自麻省理工学院(MIT)、Adobe、上海交通大学的研究者提出了一种用于压缩条件 GAN 的通用方法。这一新技术在保持视觉保真度的同时,将 pix2pix,CycleGAN 和 GauGAN 等广泛使用的条件 GAN 模型的计算量减少到 1/9~1/21。该方法适用于多种生成器架构、学习目标,配对或非配对设置。
目前该研究的论文已被 CVPR 2020 大会收录,有关 GAN 压缩框架的 PyTorch 版实现也已开源。 项目链接:https://github.com/mit-han-lab/gan-compression GAN 压缩后的具体性能如何?在研究者们展示的 Demo 中,使用 CycleGAN 为视频中的马添加斑马条纹所需的算力不到 1/16,帧数提高了三倍,而且效果还有所提升:
值得一提的是,该研究所使用的硬件平台是英伟达的边缘 AI 计算芯片 Jetson Xavier GPU。根据官方公布的数据,Jetson Xavier 的 INT8 算力为 22+10TOPS,骁龙 865 则是 15TOPS。压缩后的 GAN 现在看来已经可以跑在机器人、无人机等小型设备上了,未来塞进手机指日可待。
基于这些修改,现在可以将同一个压缩框架应用于不同类型的 cGAN。而且,与原始的非配对训练设置相比,使用上述伪对(pseudo pair)进行学习可以使训练更加稳定,并产生更好的结果。 2. 从教师判别器中学习 尽管此研究致力于压缩生成器,但判别器 D 中储存着 GAN 的有用信息。因此,研究者利用了同一个判别器架构,使用来自教师判别器的预训练权重,与压缩生成器一起微调判别器。 在实验中,研究者观察到,预训练判别器可以指导学生生成器的训练。使用一个随机初始化的判别器通常会导致训练不稳定及图像质量下降。这个 GAN 的目标可以写成以下形式:
在上式中,研究者使用来自教师判别器 D′的权重来初始化学生判别器 D。他们使用一个标准的极小化极大优化器训练 G 和 D。 3. 中间特征蒸馏 知识蒸馏是 CNN 模型压缩用到的一种普遍方法。通过匹配输出层 logit 的分布,可以将来自教师模型的暗知识(dark knowledge)迁移到学生模型中,提高学生模型的性能。然而,条件 GAN 通常会输出一个确定的图像,而不是概率分布。 为了解决上述问题,研究者匹配了教师生成器的中间表示。中间层包含的通道越多,它们所能提供的信息越丰富,学生模型所获取的输出之外的信息也越多。蒸馏目标如下:
其中,G_t(x) 和 G′_t(x) 是学生和教师模型中第 t 个被选层的中间特征激活,T 表示层数。 4. 完整优化目标 最终目标可以写成如下形式: