【论文笔记】AutoML for MCA on Mobile Devices——论文解读与代码解析

理论部分

方法介绍

本节将详细介绍AMC的算法流程。AMC旨在自动地找出每层的冗余参数。 AMC训练一个强化学习的策略，对每个卷积层会给出其action（即压缩率），然后根据压缩率进行裁枝。裁枝后，AMC算法直接对裁剪后的网络惊醒精度的测试，不经过fine-tune过程，这种精度测量结果是fine-tune后得到最优模型的精度的一中近似，而且可以加速整个训练过程。

问题定义

在深度神经网络中，通过建超每层的参数个数，达到模型压缩和加速的效果。通常有2中裁枝方法：细精度裁枝和结构化裁枝。细精度裁枝【19】直接将不重要的参数去除掉，不考虑结构信息，以单个参数（scale）为单位裁剪，这种裁枝可以做到很高的压缩率和很少的精度损失，但同时这种算法导致稀疏的不规则参数结构（无法用通常的张量或矩阵表示），需要使用特殊设计的硬件做加速（比如EIE【18】）。结构化裁枝则是粗精度的，直接去除掉一个张量中规则的区域（比如按channel裁剪，或者按整列，整行去裁剪）裁剪后的参数依然是规则的，依然可以使用张量或矩阵表示，不需要额外的硬件或软件支持。本文中将采用结构化的裁枝方案，对卷积层和全联接层的通道进行裁枝。

AMC的目标在于对每层都找出有效的稀疏度（即压缩率），在之前的工作【38，31，22】中，压缩率通常是人工决定的。以卷积层为例，原始的参数张量为 ，其中n，c分别表示输出通道数和输入通道数，k是卷积核大小。对于细精度裁枝方法，稀疏度表示为0元素个数除以所有元素的个数，而对于通道裁枝，我们直接将参数张量压缩为 ，因此稀疏度就是 .

使用强化学习自动压缩

AMC利用强化学习对action空间（在本文即压缩率）进行有效地搜索，下面将介绍本文使用的强化学习框架的细节设定。

状态空间（State Space）

对每层卷积，AMC使用11个参数来表示该层的状态 :

其中t表示第t层，其参数张量为 ；输入特征的维度为 ；FLOPs[t]表示该层的FLOPs，即该卷积层的操作数；reduced表示在之前所有层经过压缩已经减少的FLOPs；rest表示在后面所有层中剩余的FLOPs。在传入强化学习的agent之前，这些值会被归一化至【0，1】。状态 是agent区分不同卷积层的基本特征。

行为空间（Action Space）

当前大多工作中，强化学习的行为空间为离散的（比如以通道数个数为action，仅取64，128，256，512），这种粗精度的行为空间可能影像模型结构的精度，因为模型压缩后的精度，对压缩率很敏感。同时随着层数变深，离散的行为空间扩大，搜索空间变得越来越大，无法做到有效的搜索。因此，AMC提出使用连续的行为空间，以压缩率为action，则每个action： a∈[0,1]，这种搜索空间可以得到更准确的压缩率。

DDPG Agent

图1为AMC算法的整体流程。强化学习的agent接收 层的状态 ，输出为该层可以达到的稀疏度 （action），根据稀疏度取压缩后的channel个数，使用通道裁枝方法对网络进行压缩。然后将下一层 输入到agent中，对下一层可以达到的稀疏度进行估计。在完成最后一层的压缩后，直接用压缩后的模型在验证集上测试，将测试准确率作为反馈，传给agent。

AMC使用DDPG策略控制输出的连续的压缩率。DDPG在训练过程中，需要兼顾2个方面：exploration与exploitation。Exploration是将我们在确定性函数得到的行为加入噪声，以达到探索潜在更优的action，而AMC中噪声分布采用截断的正态分布（truncated normal distribution）

在exploitation阶段，噪声σ初始化为0.5，之后每个episode都以指数减少。

根据Block-QNN【54】的方法，每个episode都是从状态 ，得到行为 ，以R为反馈回传给agent，在将状态转到 ，其中R是网络压缩后的奖励。在更新阶段，为了减少梯度估计的方差，使用如下的损失函数：

为防止短期奖励比重过大，折扣因子γ设定为1。

搜索策略

限制资源消耗的压缩策略

图2 对action的约束

通过限制action（即每层的稀疏度），我们可以准确得到目标压缩率，与工作【57，4，54】类似，AMC中强化学习中使用的奖励函数为：

然而这种奖励函数中缺少对模型大小的限制，因此AMC在计算目标压缩率时限制action的解空间（即行为空间）。例如，我们要以细精度裁枝压缩模型大小，在计算前几层卷积核稀疏度时，我们使用任意的行为解（action），之后当我们发现在对所有层以大压缩率裁枝，都无法满足资源限制时，我们开始对行为解（action）进行限制。图2为整个算法的流程，这个算法中并没有指定资源限制一定是模型大小，可以由其他限制代替：比如FLOPs，或者整个网络在移动设别上一次前传的时间。作者进行了一系列实验证明，尽管强化学习的agent没有关于资源限制的刺激（AMC是通过对action作限制，而在奖励函数中没有这种资源限制），它也可以准确地给出目标压缩率。

最优精度的压缩策略

通过改变奖励函数，我们可以找到没有精度损失条件是，对压缩率的限制。作者通过实验，观察到Error与log（FLOPs）或参数的个数成负相关，基于此，作者提出以下奖励函数：

上述奖励函数不仅对Error很敏感，同时包含了操作数（FLOPs）或模型大小（#Param）的刺激。作者通过实验，证明这种奖励函数可以自动找到在保证精度的条件下，压缩率的范围。

实验部分

待完善

代码解析

• tf: https://github.com/mit-han-lab/amc-models
• pytorch: https://github.com/mit-han-lab/amc

参考资料

1. https://zhuanlan.zhihu.com/p/63299393