CVPR 2020 Oral | DMCP: 可微分的深度模型剪枝算法解读

模型轻量化是深度学习工业化中一个绕不开的话题。很多神经网络模型虽然性能很好，但是计算开销巨大，难以在手机等端上设备高效运行。因此，如何得到计算开销满足要求，同时性能也可支撑工业应用的模型一直是大家研究的重点。其中，模型剪枝就是此类方法中比较重要的一项技术。它通过裁剪神经网络的通道数来降低计算开销，同时根据特定剪枝算法来选取裁剪掉哪些通道从而降低该过程中的性能损失。在CVPR 2020上，商汤研究院被接收为oral的论文DMCP提出了一种基于马尔可夫过程的剪枝算法，为模型剪枝提供了新的思路。该工作将模型剪枝建模成了马尔科夫过程，其中的转移概率可以通过可微分的方式来进行优化，取得了非常好的效果。

动机

传统的模型剪枝方法一般分为三步（如图一所示）：训练原模型；用剪枝算法剪掉“不重要”的通道；将剪枝后模型参数进行微调。而近年来，有的工作提出剪枝后的模型参数其实并不重要，直接将剪枝模型参数初始化后重新训练，也可以达到同样的、甚至更高的精度。因此可以将模型剪枝作为模型结构搜索问题来解决，利用搜索算法搜索出模型每一层的通道数。

图一 传统的模型剪枝pipeline

先前的将模型剪枝看作模型结构搜索的工作大多都采用强化学习算法或者遗传算法，而这些算法需要在原模型中大量的采样子结构来评估精度，甚至需要花大量时间来训练子结构，因此很难泛化到更大更复杂的网络。同时，模型结构搜索问题中也存在此问题。DARTS[1]提出了一种可微分的方法来提高搜索效率。然而这种可谓分的方法不能直接套用在模型剪枝中。首先二者的搜索空间不同，其次DARTS中神经网络每一层的可选操作是互相独立的，但是模型剪枝中的选择隐含着一种互相依赖的关系，例如一层有k+1个通道的话，隐含着这层至少有k个通道。

因此，DMCP选择将剪枝的过程建模为一个马尔科夫模型。图二展示了一层通道数为5的卷积层的剪枝过程。其中S1表示保留第一个通道，S2表示保留第二个通道，以此类推。T表示剪枝完毕。概率p则为转移概率，通过可学习的参数计算得到，后文中会详细介绍。

图二 单层剪枝过程示意图。

方法

(1) 优化剪枝空间

在传统的剪枝方法中，会为每个通道计算“重要性”来决定是否保留它。而当我们把模型剪枝看作模型结构搜索问题后，不同模型的区别则在于每一层的通道数量。如果仍然每个通道单独判断，就会产生同样的结构，造成优化困难。如图三所示：情况1中，最后两个通道被剪掉，情况2中，第2个和第4个通道被剪掉，而这两种情况都会产生3个通道的卷积层，使剪枝空间远大于实际网络个数。

图三 剪枝空间冗余示意图

因此，DMCP采用保留前k个通道的方式，大大缩小了剪枝空间。

(2) 建模剪枝过程

在一层卷积层中，定义为保留前k-1个通道的情况下，保留第k个通道的概率。那么保留前k个通道的概率则为，且：

（1）

由(1)中的剪枝方式，我们可以得出，已知保留第k-1个通道后，保留前k-2个通道与保留第k个通道条件独立，且不保留第k-1个通道时，保留第k个通道的概率为0，则有下面的式子成立：

（2）

其中pk为马尔科夫模型中的转移概率。这样，通过在优化完毕后的马尔可夫模型上采样就可以得到相应的剪枝后的模型。

(3) 学习转移概率

图四 可微分马尔可夫模型示意图

图四描述了如何将马尔科夫模型中的转移概率与原模型进行结合，从而使转移概率可以通过模型损失函数的梯度来更新。该过程大致可分为三个步骤

步骤一、计算转移概率pk。

pk可以由如下公式计算得到：

（3）

其中是可学习的参数。在该方法中，每层至少保留了一个通道，因此p1 = 1。

步骤二、计算出每一个通道被保留的的边缘概率p(wk)。

（4）

由上述公式可以看出，边际概率p(wk)可以被化简为前k个转移概率的乘积，并且随着k的增大而减小。

步骤三、将每一个通道的边缘概率与对应通道的输出相乘，作为下一层的输入。

（5）

若第k个通道的边际概率趋于0，则表示该通道的输出趋于0，从而可以被剪掉。通过这种建模方式，可以保证每一层中靠后的通道先会被剪枝。

(4) 训练流程

图五 训练流程示意图

如图五所示，DMCP的训练可以分为两个阶段：训练原模型和更新马尔科夫模型。这两个阶段是交替进行来优化的。

阶段一，训练原模型。

在每一轮迭代过程中，利用马尔科夫过程采样两个随机结构，同时也采样了最大与最小的结构来保证原模型的所有参数可以充分训练。所有采样的结构都与原模型共享训练参数，因此所有子模型在任务数据集上的精度损失函数得到的梯度都会更新至原模型的参数上。

阶段二，更新马尔科夫模型

在训练原模型后，通过前文中所描述的方法将马尔科夫模型中的转移概率和原模型结合，从而可以利用梯度下降的方式更新马尔科夫模型的参数，其损失函数如下：

（6）

其中为模型在任务数据集上的精度损失函数，为计算量约束，为超参数。这里我们使用了模型期望FLOPs来约束模型的计算量（实际应用中可以方便的换用latency等指标来约束模型的计算量），其损失函数如下:

（7）

上述等式中的为目标计算量，模型的计算量期望为每一层的计算量期望的总和。

（8）

实验结果

作者在ImageNet数据集上对比了其他最新模型剪枝方法，在各种计算量下，DMCP在MobileNet-v2和ResNet上均超过现有方法，如下表所示：

传送门

DMCP代码目前已经开源，欢迎各位同学使用和交流。

论文：DMCP: Differentiable Markov Channel Pruning for Neural Networks

论文作者：Shaopeng Guo, Yujie Wang, Quanquan Li, Junjie Yan

论文地址： https://arxiv.org/pdf/2005.03354.pdf

源码地址：https:///github.com/zx55/dmcp

References

[1] Yann Le Cun, John S. Denker, and Sara A. Solla. Optimal brain damage. In Advances in Neural Information Processing Systems, pages 598–605. Morgan Kaufmann, 1990.

[2] Zhuang Liu, Mingjie Sun, Tinghui Zhou, Gao Huang, and Trevor Darrell. Rethinking the value of network pruning. arXiv preprint arXiv:1810.05270, 2018.

[3] Yihui He, Ji Lin, Zhijian Liu, Hanrui Wang, Li-Jia Li, and Song Han. Amc: Automl for model compression and ac- celeration on mobile devices. In Proceedings of the European Conference on Computer Vision (ECCV), pages 784– 800, 2018.

[4] Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. arXiv preprint arXiv:1806.09055, 2018.

[5] echun Liu, Haoyuan Mu, Xiangyu Zhang, Zichao Guo, Xin Yang, Tim Kwang-Ting Cheng, and Jian Sun. Metapruning: Meta learning for automatic neural network channel pruning. arXiv preprint arXiv:1903.10258, 2019