字节跳动开源GAN模型压缩框架，算力最高节省97.8%丨ICCV 2021

把鞋子的轮廓图：

还原成鞋子本体：

需要消耗多少计算量？

用最基本的 Pix2Pix 模型，需要消耗 56.8G；

而当 Pix2Pix 模型被一种新技术压缩后，只需要 1.219G，是原来的 1/46.6，大大节省了计算量。

这里用到的技术，就是字节跳动技术团队在计算机视觉顶会 ICCV 2021 上发表的在线多粒度蒸馏算法（Online Multi-Granularity Distillation，简称OMGD）。

这种模型压缩框架，专治 GAN 模型体积过大、太费算力，目前已经开源了代码（地址在文末）， 以及 CycleGAN 与 Pix2Pix 的预训练模型，并且已经在抖音等产品上落地。

和同类模型压缩算法相比，OMGD 这项新框架，不仅压得更小，而且压得更好。

比如在把马变成斑马的过程中：

MACs 表示消耗的计算量，括号中的数字为提升倍数

还有把夏天变成冬天：

把被分割的街景图还原成照片（注意那个骑自行车的人）：

实验表明，这项技术可以把 GAN 模型的计算量最低减少到原来的 1/46、参数量最低减少到原来的 1/82。

换算一下，就是省掉了 97.8% 的计算量。

OMGD 是怎么实现的

模型压缩通常用的是「知识蒸馏」的方法，也就是由参数多、体积大的模型作为「教师模型」，来监督优化参数少、体积小的「学生模型」，让学生模型在无需扩大体积的的情况下就能获得教师模型传授的知识。

OMGD 这项技术则是由两个分别在神经网络宽度和深度两方面互补的教师模型来优化一个学生模型，整体 Pipeline 是这样的：

该框架将不同层次的概念从中间层和输出层进行转移，可以在无鉴别器和无 Ground Truth 的设定下进行训练，实现知识的提炼，对在线蒸馏方案进行了整体优化。

在 Pix2Pix 和 CycleGAN 两大知名模型上的实验数据表明，OMGD 可以用最少的参数、最低的计算量达到最好的图像生成效果。

最右侧的 FID 分数越小，表明生成效果越好

为什么要把大模型变小？

字节跳动技术团队相关研发同学说，这是第一个用在线知识蒸馏来压缩 GAN 的技术方案，已经在抖音落地。

没错，你在抖音里看到的各种好玩的特效道具，背后都需要各种算法模型来运行，尤其是在处理图像相关的特效上，GAN 是通用的方法，这套 GAN 模型压缩方案也已经在抖音落地，比如「动态照片」：

「舞蹈手套」：

还有「三屏人生」：

不过，GAN 模型通常体量很大，需要做大量的计算才能完成，想要落地在手机尤其是算力不足的低端机上是很大的挑战。OMGD 研发团队的一位同学说：“我们会衡量模型的覆盖率，也就是一个模型有多少机型能流畅运行，成功压缩模型后可以覆盖更多手机，让更多人用上，假如原来有的模型可能要 iPhone 11 才能用，压缩之后 iPhone 7 也能用。”

所以，模型压缩是个刚需，如何让 GAN 被更多人用上、更普惠的提供服务，是技术界一直在追求的方向。

字节跳动技术团队最早 2017 年投入模型压缩方面的研究，今日头条、抖音、剪映、西瓜视频等多个 App 都有相关技术落地，还曾获得过 2020 年 IEEE 举办的低功耗计算机视觉挑战赛（LPCV）两条赛道的冠军。

在经过 OMGD 压缩之前，团队内通常使用蒸馏或剪枝等算法来完成 GAN 模型的压缩，因为模型需要的输入分辨率很大，其计算量依旧很大，没有压缩到极致。

如何才能实现更极致的压缩呢？

在研究了海量学术界现有方法之后，字节跳动技术团队同学没有找到适合公司业务的方法，转而决定自行研究，创造性地在 GAN 模型压缩上首次想到了用两个互补的教师模型来训练同一个学生模型的方式，并且取得了实验的成功。

现在，OMGD 在实践中可以比原来的方法加速 20~30%，甚至有的能达到 80%。

并且，作为可以「在线压缩」的方法，OMGD 大大降低了 GAN 模型生产的复杂性。这里的「在线」并不是我们日常生活中所指的网络在线的状态，而是指蒸馏过程一步到位地完成，“之前的 GAN 压缩方法都是分几步进行的，比如先预训练，压缩之后再训练一遍，之后还有一些其他步骤，整体比较复杂；我们的新方法只要一步就可以完成整个过程，效果也比其他方法好得多。”团队技术同学说。

现在，这类模型压缩技术不仅能节约算力和能源，更能为用户提供流畅体验，帮助创作者激发创造力，丰富生活。

相关链接

论文地址：

https://arxiv.org/abs/2108.06908

GitHub 代码及与训练模型：

https://github.com/bytedance/OMGD