AutoCompress: An Automatic DNN Structured Pruning Framework for Ultra-High Compression Rates

AutoCompress: An Automatic DNN Structured Pruning Framework for Ultra-High Compression Rates 论文阅读笔记/翻译

原文链接

摘要

结构化权重剪枝是DNN模型压缩的代表性方法，其中自动化的超参数选择是很必要的；本文提出了一种名为AutoCompress（之后简称为AC）的自动结构化剪枝框架，有以下优点：

1. 在自动化过程中高效地结合结构化剪枝方案；
2. 采用基于ADMM(Alternating Direction Methods of Multipliers) 的结构化剪枝算法，提出一种额外的权重纯化步骤；
3. 提出一种启发式，脱胎于经验导向搜索的搜索方法。

在CIFAR10和ImageNet数据集上的实验表明AC达到了权重和FLOPS的超高压缩率（120倍的参数压缩）；在智能手机上也实现了推理速度的显著提高。
本文所有模型

1. 介绍

这一部分有大量的文献引用，这里只是提一嘴其idea，标题就不列出了，想获得更多信息可阅读原文
深度网络（如VGG，ResNet）对存储和计算的高要求一直以来制约着其在移动平台和边缘设备的部署。基于此，很久以前就有人提出了以降低网络推理（也就是向前传播）期间存储和计算开销为目标的模型压缩方法。其中一个典型的方法就是权重剪枝，即在保持准确率的前提下减少权重参数的数量。

权重剪枝可分为两个分支：非结构化和结构化的。非结构化的剪枝可以任意地剪掉参数，以追求压缩率最大化，但是存储产生的稀疏矩阵需要额外的索引开销，不利于存储；结构化剪枝在保证矩阵完整性的同时减小矩阵的规模，有益于硬件加速，因此成为了近期研究的焦点。

针对深度网络的卷积层，结构化剪枝也可分为很多种，包括filter pruning，channel pruning和column pruning，近期有人提出了一种基于ADMM优化工具的系统化框架，它对结构化和非结构化剪枝都适用。

结构化剪枝策略非常灵活，涉及到如何确定每层的剪枝率。传统方法是人工选择，经过不断试错来得到较好的结果，费时费力而且导出的参数往往不是最优。受到AutoML(automated machine learning) 自动机器学习的启发，近期的工作(DRL,deep reinforcement learning)提出了自动决定每层剪枝率的方法，然而它也有自己的局限性：

1. 使用的是一种早期的，固定正则化的剪枝法；
2. 只考虑了filter pruning

作者认为自动化过程的通用流程可分为四步：

1. action sampling
2. quick action evaluation
3. decision making
4. actual pruning and result generation

本文的策略是采用ADMM剪枝算法，添加了额外的权重纯化步骤；还提出了一种经验导向的，启发式的搜索法来优化DRL。

作者在智能手机上做的实验基于他们的编译器协助的移动深度网络加速框架。

2. 相关工作

结构和非结构化剪枝

二者区别在上文已经提到了
下面给出了filter pruning，channel pruning和filter shape(column) pruning三种结构化剪枝方法的图示

filter pruning直接将一整个filter（也可叫卷积核）剪掉；channel pruning将所有filter的某一通道剪掉；filter shape pruning修改filter的形状，将所有filter的某一部分剪掉。
另外值得一提的是filter pruning和channel pruning有一定相关性：对第i层采用filter pruning时会使第i+1层的对应通道也被剪掉。

ADMM

关于admm的简明理解可以参考此处
这一部分涉及一点数学知识，没完全搞懂，大概意思是结构化剪枝问题可以转化为形如$$mi{n}_{x}f(x)+g(x)$$的优化问题，其中$f(x)$是损失函数，$g(x)$是对参数数量的约束

AutoML

NAS（神经网络架构搜索）采用了AutoML。
在传统的机器学习方法中，准确率高低很大程度上取决于特征的质量；为了减小人工筛选特征的开销，自动特征工程会生成合适的特征集。

3. AC 框架

给定一个经过预训练的深度网络，AC会决定每层的剪枝率和剪枝策略。

3.1 自动过程的通用方法

四个步骤上文已经提到了，其中的action表示对超参数选择取样；由于超参数空间很大，前两步应该快速进行以降低训练时间，作者使用简洁启发式的方法，基于被选择的超参数消除每层中的部分参数，然后评估准确率；第三步根据超参数的取样结果和评估结果确定超参数的值；第四步使用优化算法产生剪枝结果，优化算法每轮只运行一次。
重复上述步骤4-8次。

3.2 提升性能的三个根源

一是剪枝策略，作者发现filter pruning 配合 filter shape pruning 表现最好，比较结果列在附录；

二是在上述的第四步中使用 ADMM 作为优化算法，而且还添加了一个纯化步骤；具体的操作在之后详述；

三基本上是在批判DRL，启发式的搜索算法。

3.3 核心剪枝算法

算法中最主要的一步就是ADMM regularization
首先考虑一个深度网络的损失函数 f ( { W i } , { b i } ) , f(\{W_i\},\{b_i\}), f({Wi​},{bi​}),其中$W_i$和$b_i$是第i层的参数；那么剪枝问题就可以被定义成 min ⁡ { W i } , { b i } f ( { W i } , { b i } ) , \min_{\{W_i\},\{b_i\}}f(\{W_i\},\{b_i\}), {Wi​},{bi​}min​f({Wi​},{bi​}),$$s.t.W_i\in S_i,\text{for all i}，$$
$S_i$代表剪枝要求
这里不太明白为什么可以独立地考虑每一层对损失函数的影响
接下来引入指示函数
g i ( W i ) = { 0 if  W i ∈ S i , + ∞ otherwise g_i(W_i)=

$$\begin{cases} 0 & \text{if $W_i \in S_i$,}\\ +\infty & \text{otherwise}\\ \end{cases}$$ gi​(Wi​)={0+∞​if Wi​∈Si​,otherwise​
引入附加变量$Z_i$和对偶变量$U_i$，构造一个增广拉格朗日函数，接着分解成两个子问题：
min ⁡ { W i } , { b i } f ( { W i } i = 1 N , { b i } i = 1 N ) + ∑ i = 1 N ρ i 2 ∥ W i − Z i k + U i k ∥ F 2 \min_{\{W_i\},\{b_i\}}f(\{W_i\}_{i=1}^N,\{b_i\}_{i=1}^N)+\sum_{i=1}^N\frac{\rho_i}{2}\Vert W_i-Z_i^k+U_i^k\Vert_F^2 {Wi​},{bi​}min​f({Wi​}i=1N​,{bi​}i=1N​)+i=1∑N​2ρi​​∥Wi​−Zik​+Uik​∥F2​
min ⁡ { Z i } g i ( Z i ) + ∑ i = 1 N ρ i 2 ∥ W i k + 1 − Z i + U i k ∥ F 2 \min_{\{Z_i\}}g_i(Z_i)+\sum_{i=1}^N\frac{\rho_i}{2}\Vert W_i^{k+1}-Z_i+U_i^k\Vert_F^2 {Zi​}min​gi​(Zi​)+i=1∑N​2ρi​​∥Wik+1​−Zi​+Uik​∥F2​
分别求解
这部分涉及一些数学知识实在理解不能，先挖个坑
ADMM regularization的特点是动态正则化，在每轮迭代都会调整，惩罚不同的参数，这是它优于传统算法固定L1或者L2正则化的原因。

3.4 权重纯化

这一步位于基于ADMM的剪枝操作之后。
第2节已经提到，剪掉第i层的一个filter会导致第i+1层对应的一个通道也被剪掉，因此这种剪枝方法效率是比较高的。
另一方面由于ADMM regularization是一种动态的L2正则化，并且具有非凸性 （？），所以剪掉大量非零但是很小的参数仍可保持准确率。
总的来说方法就是在ADMM regularization之后设置基于filter pruning和column pruning的两个threshold值，剪掉低于threshold值的对应filter和column。

3.5 $\star$自动压缩框架

首先明确一下，这一节描述的内容应用于之前提到自动化过程的步骤一和二，即基于启发式算法对参数取样，评估，作为决定每层剪枝率和剪枝策略的依据

整个框架分为两个阶段，阶段一根据ADMM算法做剪枝（其实只是正则化？），阶段二做纯化（实际剪枝？），每个阶段都会做多轮迭代，两个阶段的执行流程是相同的，只有最后是进行ADMM还是纯化的区别。图示如下：

下面以阶段一为例来说明。
在每轮迭代的开始会将压缩率（针对参数数量或者FLOPS）设定为2，这样只经过两轮就可以达到四倍的压缩率。
每轮迭代的第一步是取样，这里使用的是基于导向搜索增强的模拟退火算法（模拟退火算法有助于脱离局部最优，可以找到近似最优解，一个简单的解释）。
具体做法是这样的：

1. 设定一个总体的剪枝率$C_t\approx 2$,
2. 将所有层按照保留参数数量排序并且为其分配一个随机的剪枝率（和基于filter和column剪枝的权重），但是要保证参数更多的层的剪枝率更高，
3. 用$C_t$标准化剪枝率
   这样就生成了一个初始的取样方案$A_t^0$，
4. 在所有层的子集上定义一个改变剪枝率和剪枝策略的perturbation（扰动），p也是满足参数越多剪枝率越大的,
5. 用p修改$A_t^0$，进行一次快速评估（即按照$A_t^0$的剪枝率和剪枝策略剪掉对应层中值最小的参数，在3.1中提到），根据退火算法修正$A_t^0$
   得到剪枝方案，执行ADMM算法。

上述流程伪代码如下：

4. 实验和讨论

作者使用Pytorch，在cifar-10上用VGG-16和ResNet-18，ImageNet上用VGG-16和ResNet-18/50做了实验，并且主要关注卷积层的剪枝，针对提升性能的三个来源分别做了实验。实验结果可查阅原文

4.1 在cifar10上的实验

得到以下结论：

1. 仅仅使用搭配纯化步骤的ADMM优化算法（根源二），作者提出的算法在精确度和剪枝率上都优于一些前辈算法（2PFPCE,NISP,AMC）；
2. 在手动决定超参数（没采用模拟退火？）的前提下，结合基于filter和column的剪枝策略在参数和FLOPS的压缩上都优于仅仅基于filter的剪枝策略；
3. 采用模拟退火后表现超过了DRL和手动决定超参数的版本。

对于另一篇论文Rethinking the value of network pruning（思考剪枝结果是得到的结构还是参数值），作者也做了相关实验，发现：
仅使用基于filter的剪枝策略，放弃已有参数值从头开始训练可以恢复原先的精确度
对此作者认为基于filter和channel的剪枝策略本质都是在寻找一个更小的网络结构；以这个角度看，作者的AC框架确定每层的压缩率，也做了同样的事情。
但是另一组实验表明基于混合剪枝策略，从头训练不能恢复到原来的精确度，作者认为由于混合剪枝策略不光产生了更小的网络，同时也调整了filter的形状，这就需要原先大网络的一些细节了。（感觉filter的形状与原本网络的参数值并没有直接的关系，作者的解释有些模糊）

4.2 在ImageNet上的实验

结果与cifar10上类似，具体信息查看原文表格

5 结论

对前面内容的概括