pytorch基于卷积层通道剪枝的方法

本文基于文章：Pruning Convolutional Neural Networks for Resource Efficient Inference

https://arxiv.org/abs/1611.06440

下面以ResNet模型为例，讲解剪枝方法。

 

1  评估神经元（即通道）重要性

1.1 给模型中所有conv layer重新编号排次序；

1.2 计算每个conv layer所有通道rank值，并按通道归一化；

1.3 统计模型所有conv layer的通道数；

1.4 根据设置拟剪枝数量num（默认512），对所有convlayer通道rank值按从大到小排序，即确定rank值最小的num个通道，记录并返回它们的卷积层号和通道编号。

2  去除不重要神经元

2.1 根据拟剪枝列表（即卷积层号和通道编号），在训练好或加载的模型中去除；

2.2 如何去除，即把conv layer中按2.1中的列表，去除通道a，使a的前后通道相连，使第一卷积层输出通道数与第二层输入通道数相同，最后一个卷积层输出需与第一个全连层        输入通道数相同，这里全连层不计算rank值，拟剪枝的卷积层，需要conv继承操作；

       注意：残差块中需要处理公式（1）F(X)=f(x)+x中f(x)与x不同维数情况；即在convlayer通道数改变后，添加residual = out即可。这个修改虽与公式（1）本意有些出入，但通过实验发现，在模型大小减少一半的情况下，仍可以较好地与原模型精度差在1%以内；

2.3 BatchNorm层通道数修改

       当所有卷积层剪枝结束，依据邻近上一个卷积层输出通道数，通过BatchNorm层继承方式，它需修改成同样的通道数即可；

3 剪枝模型训练

       全网参数更新，也可设置微调特定层参数更新

4 判断剪枝是否结束，若停止到下一步微调全网参数，否则重复1---3步骤；

5微调模型

       使用与训练同一个数据库，也可用另外数据库，一般来说学习率要小于训练的1到2个数量级，微调结束后保存剪枝模型。

6 问题与讨论

6.1 具体模型需定制剪枝，需依照其模型生成文件，也包括forward函数，这要在计算Loss时使用；

6.2 测试阶段

      需保证剪枝模型可forward操作，设置好Variable参数，即volatile=True，以免GPU资源泄露；

 

具体代码参见：https://github.com/eeric/channel_prune

2019.10.15日添加：从实践来看，模型剪枝算法存在自身缺陷，想要获得模型小速度快，精度有提高的办法是使用知识蒸馏训练。