『量子位』GPU、CPU统统能加速，朱俊彦团队提出GAN压缩算法：计算量减20倍( 二 ) 边策鱼羊发自凹非寺量子位报道|公众号Q

文章图片
最后完整的损失函数为：

文章图片
其中超参数λ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
Python3
CPUorNVIDIAGPU+CUDACuDNN
同时，研究团队也准备了CycleGAN和pix2pix两个版本的PyTorchColab ，在线即可体验。
试试用CycleGAN把普通马变斑马的效果。
首先把GANCompression这个项目克隆下来：
!gitclonehttps://github.com/mit-han-lab/gan-compression.git
下载测试数据集：
!bashdatasets/download_cyclegan_dataset.shhorse2zebra
下载预训练模型：