理论部分
方法介绍
本节将详细介绍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。
搜索策略
限制资源消耗的压缩策略
通过限制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
参考资料
- https://zhuanlan.zhihu.com/p/63299393