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

字节跳动开源GAN模型压缩框架节省高达97.8%算力

时间:2023-03-14 21:46:07 科技观察

还原鞋子轮廓:需要消耗多少算力?最基本的Pix2Pix模型,需要消耗56.8G;而当Pix2Pix模型经过新技术压缩后,只需要1.219G,是原来的1/46.6,大大节省了计算量。这里使用的技术是字节跳动技术团队在计算机视觉顶级会议ICCV2021上发表的在线多粒度蒸馏算法(OnlineMulti-GranularityDistillation,简称OMGD)。这个模型压缩框架是专门为处理GAN而设计的模型太大且计算量大。目前代码已经开源(地址在文末),以及CycleGAN和Pix2Pix的预训练模型,已经在抖音等产品中实现。与同类模型压缩算法相比,OMGD新框架不仅压缩得更小,而且压缩得更好。比如在马变成斑马的过程中:MACs代表消耗的计算量,括号中的数字是改进的倍数;以及把夏天变成冬天:将分割后的街景图像还原成一张照片(注意骑自行车的人):实验表明该技术可以将GAN模型的计算量减少到最少的1/46,参数数量也减少到最少的1/82。换句话说,97.8%的计算都被保存了。OMGD是如何实现的模型压缩通常采用“知识蒸馏”的方法,即用参数多、体积大的模型作为“教师模型”,对参数少、体积小的“学生模型”进行监督优化,使学生模型可以在不扩大体积的情况下获取教师模型传授的知识。OMGD是一种通过在神经网络宽度和深度方面相互补充的两个教师模型来优化学生模型的技术。整体的Pipeline是这样的:框架从中间层到输出层传递不同层级的概念。可以在无判别器、无GroundTruth的设置下进行训练,实现知识的蒸馏,并对在线蒸馏方案进行整体优化。在Pix2Pix和CycleGAN这两个著名模型上的实验数据表明,OMGD能够以最少的参数和最低的计算量达到最好的图像生成效果。最右边的FID分数越小,生成效果越好。为什么要把大模型变小?字节跳动技术团队的相关研发同学表示,这是第一个通过在线知识蒸馏来压缩GAN的技术方案,已经在抖音实现。没错,你在抖音看到的各种有趣的特效道具,背后都需要运行各种算法模型,尤其是处理图像相关的特效,GAN是一种通用的方法,而这套GAN模型压缩方案也有在抖音中已经实现,比如“动态照片”:“跳舞的手套”:和“三屏生活”:但是GAN模型通常非常庞大,需要大量的计算才能完成。手机,尤其是计算能力不足的低端机,是一个很大的挑战。OMGD研发团队的一位同学表示:“我们会衡量模型的覆盖率,即一个模型有多少个模型可以流畅运行。成功压缩模型后,可以覆盖更多的手机,让更多的人使用它。如果有iPhone11可能需要使用该型号,iPhone7压缩后也可以使用。因此,模型压缩是刚需,如何让GAN更容易被更多人使用,提供更普惠的服务,是技术社区一直追求的方向。字节跳动技术团队在2017年首次投入模型压缩研究。今天的头条、抖音、简影、西瓜视频等APP都实现了相关技术,还获得了2020年IEEE低功耗计算机竞赛的冠军,VisionChallenge(LPCV)两个赛道的冠军。在OMGD压缩之前,团队通常使用蒸馏或者剪枝等算法来完成GAN模型的压缩,因为模型需要很大的输入分辨率,其计算量还是很大的,还没有压缩到极致,如何才能做到更极致的压缩?字节跳动技术团队的同学在研究了大量学术界已有的方法后,没有找到适合公司业务的方法,于是他们决定自己做研究,并创造性地想到首次使用两个互补的教师模型来训练同一个学生模型的GAN模型压缩方式,并取得了实验成功。现在OMGD在实践中可以比原来的方法提速20~30%,有的甚至可以达到80%。而且,作为一种可以“在线压缩”的方法,OMGD大大降低了GAN模型制作的复杂度。这里的“在线”并不是指我们日常生活中网络的在线状态,而是指一步完成的蒸馏过程。“之前的GAN压缩方法都是分几个步骤进行的,比如先进行预训练,压缩之后再训练,之后还有一些其他的步骤,整体比较复杂;我们的新方法可以在一步到位,效果比其他方法好太多了。”团队的技术同学表示,现在这种模型压缩技术不仅可以节省计算力和能源,还可以为用户提供流畅的体验,帮助创作者激发创造力,丰富生活。相关链接论文地址:https://arxiv.org/abs/2108.06908GitHub代码和训练模型:https://github.com/bytedance/OMGD