AutoGAN-Distiller: Searching to Compress Generative Adversarial Networks, ICML2020

                                   

通过NAS对GAN网络进行模型压缩的一篇文章，作者用图像翻译的CycleGAN和超分的ESRGAN来示例，这俩一个是17年一个是18年的模型。

原文章相关资料如下：

Paper：https://arxiv.org/abs/2006.08198

Github：https://github.com/TAMU-VITA/AGD

 

网络结构的总体设计

总体结构上，很多NAS模型都采用了有向无环图（DAG）型搜索空间，但是这种结构会带来大量不规则的密集连接，计算量大，而且文章提到这种结构对硬件的利用不够友好——因为会阻碍数据的并行处理，带来latency。

因此作者没选DAG，而是选择顺序的搜索空间，这样每个单独的搜索模块（也叫节点）只包含一个运算操作，这些节点顺序连接，没有分支。

NAS搜索由搜索空间和代理任务两部分构成，分别对应平常任务中的网络结构设计和损失函数设计与优化两部分。

 

搜索空间

搜索对象：操作类型（如1*1卷积）与channel宽度。

 

                       

                         

 

上面两幅图分别是图像翻译和超分的改进网络结构示意图，其实就是固定住Stem和Header（特征提取与复原部分），去搜索虚线框中的结构（特征处理部分）。

主要看超分这边，为了提高效率，作者将多数计算都放在了在低分辨率的特征空间中。对于虚线框中的模块，本来是ESRGAN在SRGAN的基础上将残差与DenseNet结合得到的RRDB，特点是连接紧密、更容易训练，但是计算量大、影响数据并行，于是就在对这部分搜索操作类型和channel宽度，进行替换。另外作者也没用BN层，是为了避免产生伪影。

操作类型方面，其搜索范围是以下几个，都是较流行的模块：

Conv 1*1、Conv 3*3

Residual Block-ResBlock：Conv 3*3带有skip connection

Depthwise Block-DwsBlock：Conv1*1 + Depth-wiseConv 3*3 + Conv 1*1带有skip connection

 

代理任务

                                          

代理任务对应平常训练中的损失函数及优化部分，本文的损失函数如（1）式，G0为预训练模型（teacher），G为要训练的生成器，d指两者输出的距离，F指的是计算cost，α、γ是控制网络结构的参数，通过softmax确定每一层的操作类型和信道宽度。

loss很自然地体现两个思想：

（1）压缩后的模型输出和原来相差无几 。

（2）减少计算负担。

具体来说，前一项损失d(G,G0) 包括了content loss（为了避免color shifting）、perceptual loss（保证视觉效果）以及total variation loss（保证元素级相似性）三项，不算新鲜；而后一项loss衡量computional budget则直接用的FLOPs（浮点数计算量）。

可以发现压缩过程与判别器无关，作者通过实验证明了就算优化时带上判别器效果也没有提升，反倒让训练过程不稳定了，这个经验可以吸取。个人也是认为在增强任务中GAN训练的目的就是在G与D的博弈下得到生成器，判别器只是配角。

 

训练过程

文章的伪代码（图比较长，放在下面）对训练过程讲的比较清楚，分为三步：

（1）预训练：用一半训练集预训练网络参数。

（2）搜索：交替更新网络参数和结构参数（α、γ），搜索网络结构。

（3）重新训练：选定每层的操作类型与channel宽度，重新初始化网络参数后再训练一次。

                                                            

压缩模型的效果

                                        

在最终效果方面，（1）评估计算负担是用256*256的图像进行4倍上采样实验，可见AGD压缩后的模型所需算力及内存都显著减少，（2）同时能保证PSNR下降不多，（3）视觉效果基本看不出差别。

PSNR评价质量时用的是超分任务常用的四个数据集，略去了manga109，因为五个数据集中只有它是手绘漫画数据集，其他四个都是照片、自然图像，gap比较大且应用性没那么强，略去它是可以理解的。

最后放一下文章中的视觉效果对比图，压缩后的模型视觉效果基本和原来一致。