【论文笔记】AMC：AutoML for Model Compression and Acceleration on Mobile Devices
论文：https://arxiv.org/abs/1802.03494
代码：https://github.com/mit-han-lab/amc-release
简介
本篇论文是MIT韩松老师组提出的一种模型压缩方法，其核心思想是使用强化学习技术来实现自动化压缩模型。目前对于领域问题的解决方案通常采取人工压缩模型，手工压缩方法需要相关领域的专家知识，需要人工在推理速度，大小与准确率之间权衡，因此人工压缩难以得到最优的压缩策略。与人工设计模型相比使用强化学习的产生压缩策略从数据出发能够找到更合适的压缩策略。该篇文章讲FLOPs减少4倍的情况下，对VGG-16模型压缩，比人工压缩策略策略高2.7％的精度。并将这种方法应用到更适合移动端设备的Mobilenet模型上，使模型更加轻量化。
相关工作
与现有的技术相比 （NAS nerual architecture search [57]、NT network Transformation[6]、N2N network to network[2] ）本文提出的方法AMC不是对网络结构的路径搜索，而是采用强化学习中的DDPG（深度确定性策略梯度法）来产生连续空间上的具体压缩比率，通过约束FLOPs和准确率设定奖励值，因此得到的压缩策略是适应当前硬件环境的。
方法
方法整体结构图如下
图 1：AMC 方法图示。左：AMC 替代人工使模型压缩完全自动化，而且表现优于人类。右：将 AMC 映射到强化学习问题场景。珠层压缩 一个pretrained的 MobileNet。Agent（DDPG）从 t 层接收嵌入 s t s_t st，输出稀疏比率 a t a_t at。用 a t a_t at 对层进行压缩后，智能体移动到下一层 L t + 1 L_{t+1} Lt+1。评估所有层都被压缩的修剪后模型的准确率。最后，奖励 R 作为准确率和 FLOPs 的函数被返回给Agent。
本文使用DDPG方法在前问题场景下，agent的动作空间映射为 待压缩模型每一层的压缩率 ，本文的目标是保持准确率的同时减少FLOPS与压缩模型的大小，因此回报函数可以根据这些指标进行设计，论文中提到了两种回报函数，基于FLOPs和参数量的，如下所示
S t a t e S p a c e : ( t , n , e , h , w , s t r i d e , k , F L O P s [ t ] , r e d u c e d , r e s t , a t − 1 ) State Space: (t,n,e,h,w,stride,k,FLOPs[t],reduced,rest,a_{t-1}) StateSpace:(t,n,e,h,w,stride,k,FLOPs[t],reduced,rest,at−1)
L o s s = 1 N ∑ i ( y i − Q ( s i , a i ∣ θ Q ) ) 2 Loss = \frac{1}{N} \sum_i{(y_i - Q(s_i,a_i| \theta^Q))^2} Loss=N1∑i(yi−Q(si,ai∣θQ))2
y i = r i − b + γ Q ( s i + 1 , μ ( s i + 1 ) ∣ θ Q ) y_i = r_i - b+\gamma Q(s_{i+1},\mu(s_{i+1})|\theta ^ Q) yi=ri−b+γQ(si+1,μ(si+1)∣θQ)
R F L O P s = − E r r o r ⋅ l o g ( F L O P s ) R_{FLOPs} = - Error · log(FLOPs) RFLOPs=−Error⋅log(FLOPs)
R P a r a m = − E r r o r ⋅ l o g ( P a r a m ) R_{Param} = - Error · log(Param) RParam=−Error⋅log(Param)
算法的伪代码如图所示
算法流程 主要通过DDPG算法为每一层依次产生保留率，源码中设定的动作边界为[0.2，1]，然后根据当前层原始通道数量来计算出需要保留channels的数量n，然后使用ranking的办法为通道排序然后 取出前n个通道索引作为新model保留的通道，然后分别使用 W n e w W_{new} Wnew 和 W a l l W_{all} Wall 两组权重计算 top1 acc FLOPs等指标来计算奖励值。
实验
论文分别使用CIFAR10和ImageNet数据集分别进行实验来验证，主要压缩了Resnet、Vgg和mobilenet验证，手工压缩和AMC方法调优之后模型的压缩率和准确率，实验结果AMC方法在各个维度上全面超越了手工调参的模型的指标
实验结果提出了一个很有趣的结论，作者根据实验结果提出3x3大小的权重包含了更多地冗余信息因此在3x3层给出的压缩率较高，而1x1大小的权重包含的冗余信息小因此给出的压缩率较低。这里我认为应该有更深层次的原因值得研究。
TIPs
总结一下复现过程中的一些坑，这篇论文的代码已经开源了，但是开源项目中目前只保留了mobilenet模型的剪枝源码。
1、官方开源代码在进行通道裁剪过程中 采用的通道ranking的方式是直接将每个通道上的所有权值求和，似乎有点不科学，因此我尝试使用更科学的方式（几何中位数距离，范数距离，余弦相似度等）替换ranking的方式，结果居然都没有直接求和的效果好让我很是失落。
2、同时这部分代码仅能够对mobilenet 的1x1 或 fc 层进行权重重构 reconstruction，其他的大卷积核的层写了个 Todo 留了个小坑。
3、当需要压缩其他模型时需要注意模型的多分支结构，例如resnet inception之类的，需要保持裁剪通道数的一致。另外在构建模型的时候一定要写成基于配置参数构建model的形式，不然每剪枝一次产生的model都不一样就需要手动重写新的model，会要老命的！！
最后我发现了一个Intel的蒸馏工具包
https://github.com/NervanaSystems/distiller
这份代码整合了时下比较经典的压缩剪枝算法，当然也整合了本篇论文的AMC算法。设计模式思想体现的淋漓尽致，不愧为工业界的正经代码。拜读源码好几天，真的是被秀到了。想要用起这份代码还是需要一定功力的，如果只是paper做实验推荐官方开源的代码，在基础上进行修改还是比较容易的，因为功能比较单一注释也很友好通俗易懂。