模型压缩 | 结构性剪枝Data-Driven Sparse Structure Selection 以及实际剪枝实现流程

前言

剪枝的文章很多，目前稀疏化的方法有主流的：group lasso; 控制网络结构收缩constrain the structure scale；对网络结构正则化处理regularizing multiple DNN structures。本文主要侧重于channel prune 乃至group/block等结构上稀疏化的方法。这里主要是基于图森的文章： Data-Driven Sparse Structure Selection for Deep Neural Networks。此外，如题，模型压缩是很实用的一个方向，言外之意就是工程意义很大。但是论文往往不会提及如何实际上将“繁枝”剪掉。提醒一句，就算乘以0也算FLOPs。只有要剪掉的枝从搭建的网络中消失才算哦。所以博主最后会将github代码也附上。另外，水平一般能力有限，欢迎大家斧正。 （更多模型压缩文章如各种知识蒸馏方法与实现，以后会从笔记本上整理陆续放出）。

一.文献解读

I. Networks Slimming

此文思路还是很清楚的，我想从相关文献说起。Networks Slimming-Learning Efficient Convolutional Networks through Network Slimming。
该论文pytorch实现代码地址：https://github.com/mengrang/Slimming-pytorch
思路归纳为两点：
1. 利用BN中的缩放因子gamma 作为评价上层输出贡献大小（下层输入）的因子，即gamma越小，所对应的神经元越不重要，就可以裁剪掉。

2. 利用L1正则化gamma，使其稀疏化，这样就可以在训练过程中自动评价神经元贡献大小，为0的银子可以安全剪掉。

这个思路还是很巧妙的。在这里我不得不佩服作者对深度学习的思考深度。但是，这个方法工程上会有一个需要注意的地方。根据BN原理，需要训练的因子有两种，scale 和shift。gamma是相乘因子，之后还要加上一个shift，beta.那么，gamma很小的时候，beta是不是很大？一定要保留beta.

II. Sparse Structure Selection

下面介绍图森的SSS。这里没有使用BN带的gamma，而是新增加一个参数作为衡量神经元输出，每一group/block输出重要性的因子。原理很清晰，就是在每个神经元输出上，每个残差结构卷积分支上以及一些group卷积的非identity group上之输出乘上一个缩放因子。训练过程中对这些因子施加L1正则化使得其稀疏化。这样，为0的因子对应的神经元/group/block自然而然的认为对结果不重要，安全移除。
文章的思路是非常棒的，但是我认为解决L1正则化不可导的思路也很好。首先我们知道，对于不可导的L1 norm 一般使用Proximal Gradient来反传。（这里插一句，强化学习中的PPO，近端策略优化也是解决不可导问题，第一个P的含义）
根据推导过程发现，我们要想使用该方法，显然要进行两次前传，代价太大。（原因就是要对重要性因子之差再进行一次然后根据牛顿动量优化器原理发现）唉，其实这个APG算法只不过比牛顿动量优化算法多出一个投影函数（PG算法中的概念，不懂的可以网上搜资料补补），而根据L1正则化Proximal Gradient原理，这个投影函数就是软阈值函数（有数学证明，这里就不写了，因为我也没证明）。所以，只需要在牛顿动量的优化器中对缩放因子施加软阈值函数即可，软阈值两个参数是学习率和对应缩放因子值。

二. 剪枝的实现

至此，文章的思路就介绍完毕了。源码是MXNET写的。但是，这里我们发现一个问题。正如博主开头说的，要想达到模型参数量减少目的，这只是一个开端。因为即使相乘参数为0，不把这一神经元/group/block从网络结构中拿掉，计算量不变。这是所有论文里不给说的部分，从工程上讲，这才是实际的剪枝。下面介绍一下Tensorflow如何实现剪枝。github代码也附上

1.前文加入缩放因子的网络已经训练完毕，稀疏化完毕。从ckpt中取出重要性因子的值，以及各神经元权重。

2.用numpy数组操作，记录小于设定阈值的元素位置，找到对应神经元数组位置，直接按channel整个切除。

3.参考以上切除之后的对应的权值数组搭建网络，以此设定张量大小。注意，block级别的稀疏，直接让输出跨过这一block即可。然而，channel级别的稀疏，由于残差结构的存在，相加位置张量channel要一致，是个坑。解决方案是不要对block最后一层的bottleneck稀疏即可。

4.应用指定变量初始化方法对重新定义的网络结构初始化，fineturn一下。

水平有限，欢迎大家斧正

HongYuSuiXinLang
Zhejiang University
Email: r_meng@zju.edu.cn