朱俊彦团队提出GAN压缩算法：计算量减小20倍，生成效果不变，GPU、CPU通通能加速...
边策 鱼羊 发自 凹非寺
量子位 报道 | 公众号 QbitAI
现现在，GAN的效果已经愈来愈出神入化。python
好比英伟达的GauGAN，就如神笔马良，可以凭空造物：git
不过，从无化有背后，计算量也至关惊人。github
以GauGAN为例，与MobileNet-v3这样的识别CNN相比，参数只相差1个数量级（4.2 : 93）计算强度却高出了2个数量级（0.5 : 281）。算法
如此一来，交互式部署就变得很困难。bash
更直白来讲，想要把模型部署到手机、平板这样的终端，换装变脸转性别，这些边缘设备——吃不太消。网络
为了解决这一问题，来自MIT、Adobe研究院和上海交通大学的团队琢磨出了一个通用压缩框架。架构
有多强？框架
一举将CycleGAN的计算量减小了20倍以上，将GauGAN的计算量减小了9倍，简笔画到实物图的pix2pix也能驾驭，效果却未差分毫。机器学习
值得一提的是。论文已入选CVPR 2020，代码也已开源。ide
做者团队也星光璀璨，一做是来自上海交通大学ACM班的本科生李沐阳，CycleGAN的做者朱俊彦则亲身参与、指导。
GAN压缩算法
压缩GAN，并不像压缩其余CNN模型那样容易。
主要缘由：GAN的训练动力学高度不稳定，生成器与鉴别器之间存在巨大的结构差别，让咱们很难使用现有的CNN压缩算法。
为了解决这个问题，做者提出了针对有效生成模型量身定制的训练方案，并经过神经架构搜索（NAS）进一步提升压缩率。
GAN压缩的框架以下图所示，主要分为3个部分：
一、给定一个预训练的teacher生成器G’，经过蒸馏的方法获取一个较小的“once-for-all”的student生成器G，其中包括经过权重共享的全部可能通道数。在每一个训练步骤中选择不一样的通道数ck。
二、从“once-for-all”生成器中提取许多子生成器，并评估其性能，无需从新训练，这也是被叫作once-for-all（一劳永逸）的缘由。
三、最后，根据给定的压缩率目标和性能目标，选择最佳子生成器，进行微调，并得到最终的压缩模型。
构造合适的损失函数
由于GAN压缩算法要面对CycleGAN、pix2pix还有GauGAN，这些模型的损失函数都不尽相同，因此须要构造一个新的损失函数。
统一未配对和配对学习
有些GAN是经过配对数据集学习的，有些则是非配对数据集。所以要在损失函数中加入第二项，统一非配对和配对学习的损失：
继承teacher鉴别器
尽管目标是压缩生成器，可是鉴别器会存储学习到GAN知识，由于它会发现当前生成器的缺点。
所以，咱们采用相同的鉴别器架构，使用teacher预训练的权重，并与咱们的压缩生成器一块儿对鉴别器进行微调。预训练的鉴别器能够指导student生成器的训练。
初始化的student鉴别器D使用来自teacher鉴别器D’的权重。
中间特征蒸馏
蒸馏是CNN中普遍使用的模型压缩方法。
CNN模型压缩的一种普遍使用的方法是知识蒸馏。经过匹配输出层的logits，能够将知识从teacher模型转移到student模型，从而提升后者的表现。
最后完整的损失函数为：
其中超参数λrecon和λdistill控制后两项的权重。
高效的生成器设计空间
选择设计良好的student体系结构对于最终进行知识蒸馏相当重要。
简单地缩小teacher模型的通道数并不能产生紧凑的student模型，一旦计算量的压缩比超过4倍，性能就会开始显著降低。
预测可能的缘由之一是，现有的生成器采用的图像识别模型，可能不是图像合成任务的最佳选择。
下面，做者展现了如何从现有的cGAN生成器中得到更好的架构设计空间，并在该空间内执行神经架构搜索（NAS）。
卷积分解和层敏感性
近来高效的CNN设计，普遍采用了卷积分解的方法，证实了在性能与计算之间的权衡取舍。做者发现使用分解的卷积，也有利于cGAN中的生成器设计。
然而实验代表，将分解直接应用于全部卷积层，将大大下降图像质量。可是能够只有某些层分解会下降性能，而某些层则更鲁棒。
在ResNet生成器中，resBlock层消耗了大部分模型参数和计算成本，而几乎不受分解的影响。相反，上采样层的参数要少得多，可是对模型压缩至关敏感，适度的压缩也会致使FID大幅降低。
NAS自动裁剪通道
如今的生成器在全部层上都使用手动设计，所以通道数会有冗余。为了进一步压缩模型，做者使用通道修剪自动选择生成器的通道宽度减小冗余，能够二次减小计算量。
对于每一个卷积层，能够从8的倍数中选择通道数，能够平衡MAC和硬件并行性。
解耦训练和结构搜索
为了解决该问题，做者遵循one-shot神经体系结构搜索方法的最新工做，将模型训练与体系结构搜索分离 。
先训练一个支持全部通道的“once-for-all”网络，具备不一样数量通道的每一个子网络都通过同等训练，能够独立运行，子网络与“once-for-all”网络共享权重。
在训练了“once-for-all”网络后，经过直接在验证集上评估每一个候选子网络的性能来找到最佳子网。因为“once-for-all”网络通过权重共享的全面训练，所以无需进行微调。
经过这种方式，咱们能够将训练和搜索生成器体系结构分离开来：只须要训练一次，在无需进一步训练的状况下评估全部可能的通道配置，并选择最佳的做为搜索结果。
实验结果
最终实验结果以下：
GAN压缩将乘法累加运算（MAC）的数量减了9~21倍。
这种方法将CycleGAN的计算量减小了20倍以上，将pix2pix的计算量减小了9倍，将GauGAN的计算量减小了9倍。
并且不只能在GPU上加速，在各类各类各样的CPU上也能够实现加速，并且加速效果在CPU上更明显。
代码开源
如今，这一GAN压缩框架已经开源，数据集和预训练模型都可下载。
若是想要上手亲测，须要的环境是：
Linux
Python 3
CPU or NVIDIA GPU + CUDA CuDNN
同时，研究团队也准备了CycleGAN和pix2pix两个版本的PyTorch Colab，在线便可体验。
试试用CycleGAN把普通马变斑马的效果。
首先把GAN Compression这个项目克隆下来：
!git clone https://github.com/mit-han-lab/gan-compression.git
下载测试数据集：
!bash datasets/download_cyclegan_dataset.sh horse2zebra
下载预训练模型：
python scripts/download_model.py —model pix2pix —task edges2shoes-r —stage full python scripts/download_model.py —model pix2pix —task edges2shoes-r —stage compressed
下面，就是见证换皮的时刻：
效果差距大不大，肉眼看了不算，仍是要计算一下FID（用于评估GAN生成图像的质量，分数越低表明与真实图像越类似）。
该项目提供了几个数据集的真实统计信息：
bash ./datasets/download_real_stat.sh horse2zebra A bash ./datasets/download_real_stat.sh horse2zebra B
测试的结果显示，原始CycleGAN的FID是65.687，压缩后，模型的FID是65.312，二者差距不大。
但在MAC、参数规模和延迟方面，压缩模型都要远远小于原始模型。
上海交大本科生一做，朱俊彦&韩松加持
论文一做，是上海交通大学ACM班大四本科生李沐阳。
2019年7月-今年1月，李沐阳师从MIT韩松教授和朱俊彦教授，在MIT Han Lab担任研究助理，这篇CVPR 2020论文就是在此期间产出。
目前，李沐阳的研究兴趣在于机器学习、系统以及计算机图形学等领域，他计划在今年毕业以后继续读博深造。
科研之余，李沐阳同窗的一大爱好是唱歌。在他的我的主页上，他还分享了本身的唱吧连接，以及我的MV。
论文的另外几位做者，分别是：
Ji Lin，本科毕业于清华大学，现于MIT就读博士。
丁尧尧，一样是来自上海交大ACM班的本科生，和李沐阳同在MIT Han Lab担任研究助理，受韩松教授指导。
Zhijian Liu，本科毕业于上海交大，现于韩松教授门下就读博士。
朱俊彦，李沐阳的指导者之一。这位青年大牛无需多介绍，CycleGAN做者，国际顶会ACM SIGGRAPH 2018最佳博士论文奖得到者。现为Adobe研究科学家，今年秋天将回归母校CMU担任助理教授。
韩松，MIT EECS助理教授，一样是AI业界大牛。博士毕业于斯坦福大学，曾斩获ICLR 2016最佳论文、FPGA 2017最佳论文。
传送门
论文地址：
https://arxiv.org/abs/2003.08936
GitHub：
https://github.com/mit-han-lab/gan-compression
李沐阳我的主页：
https://lmxyy.me
Colab：
https://colab.research.google.com/github/mit-han-lab/gan-compression/blob/master/cycle_gan.ipynb
https://colab.research.google.com/github/mit-han-lab/gan-compression/blob/master/pix2pix.ipynb
做者系网易新闻·网易号“各有态度”签约做者
— 完 —
3月25日晚20:00，地平线BPU算法负责人罗恒，将分享地平线BPU的设计与演进、MLPerf基准测试的理想性与局限性、提出当面AI芯片面临的机遇与挑战
戳下方二维码便可报名、加交流群~
今晚直播 | 地平线BPU从新定义极致效能
量子位 QbitAI · 头条号签约做者
վ'ᴗ' ի 追踪AI技术和产品新动态
喜欢就点「在看」吧 !