自定义代码进行Network Slim剪枝

最近的项目需要最后进行检测算法的部署，部署后的算法推理速度不是很快，需要对模型进行剪枝。

我参考的剪枝算法是2017 ICCV上的一篇经典剪枝算法：《Learning Efficient Convolutional Networks through Network Slimming》。算法的原理也很简单，需要用到的前置知识就是先理解Batch Normalization的原理，而针对BN的原理我在这篇文章Batch Normalization和梯度消失以及梯度爆炸的原理中进行了简要分析，需要的可以看看。

由于我的检测算法搭建在mmdet这个框架上，试了一下微软的开源剪枝工具nni，但是没有成功，最后打算自己写代码进行剪枝，以下是我剪枝的一个思路和具体代码实现。

1.1剪枝思路

根据以上参考论文提到的原理，我们需要对BN层进行剪枝。拿resnet50为例，我们在搭建网络的时候，一个Conv层后面一般都要接一个BN层，而BN层后面除了激活函数层外，还要连接下一个Conv层，为了方便，我们把BN层前面的卷积层记为Conv1，后面的记为Conv2。他们之间的关系为：Conv1的输出通道数和BN层的参数维度一致，也要和Conv2的输入通道数保持一致，如下图所示
在这里插入图片描述
这里拿resnet50的其中一层举例，Conv1的输入通道数为256，卷积核尺寸为[1,1]，输出通道数为64；Conv1的输出作为BN层的输入，经过BN层处理后的输出和其输入的通道数保持一致，都为64；接着将BN层的输出作为Conv2的输入，Conv2会接受64通道的输入，输出同样为64通道。

接下来介绍一下BN的剪枝方法，也很简单。BN层会针对输入的每一个通道学习两个参数，分别为 $\beta$ 和 $\gamma$ ，对应如下公式：
$y_{1} \leftarrow \gamma_1 \hat{x_{1}}+\beta_1 \equiv B N_{\gamma_1, \beta_1}\left(x_{1}\right)$
所以参数 $\gamma$ 可以看做衡量每一个通道重要性的权重。我们可以通过设定一个阈值，低于这个阈值就删掉这个通道，高于则保留。如下示意图所示。
在这里插入图片描述

1.2代码实现思路

清楚了原理接下来就是自己写代码实现这一思路。我们要对一个BN层进行剪枝，就要把与该BN层直接连接的前后两个Conv层同时进行剪枝。

首先要将整个模型加载进来，我这里使用的方法是直接用torch.load()函数：

# 保存整个网络
torch.save(model, PATH)
# 加载整个模型
torch.load(PATH)

这样加载出的模型打印出来是这样的：
在这里插入图片描述
可以看到这是模型的一个结构（我只取了resnet50的一部分），接着我们打印出它的参数来看看：

以上是该model的一部分参数维度。可以看到BN层的参数是64维的，与上一个Conv层的输出通道数对应，并且BN层有weight，bias，mean，var这4个带有维度的参数，我们要对他们进行同时修剪。

前面原理部分提到，我们是拿 $\gamma$ 参数作为衡量通道重要性的权重的，对应参数部分的weight。所以我们在对BN层进行剪枝时，首先要设定一个阈值，然后比较weight每一个值与该阈值的大小，得到大于该阈值的索引：

def find_indice(module, thresh):  #module就是一个BN层
    gamma = module.weight.data
    mask = gamma > thresh
    indices = torch.nonzero(mask).view(-1)
    return indices

接着我们利用该索引对BN层的4个参数进行修剪。除了参数的修剪，还要特别注意对结构也进行修剪，即把model上对应BN层的通道数修改成剪掉后的维度：

#对参数进行修剪
m.weight.data = m.weight.data[bn_dict["backbone.layer1.0.bn1"]]  #gamma
m.bias.data = m.bias.data[bn_dict["backbone.layer1.0.bn1"]]      #beta
m.running_mean.data = m.running_mean.data[bn_dict["backbone.layer1.0.bn1"]]
m.running_var.data = m.running_var.data[bn_dict["backbone.layer1.0.bn1"]]
#对结构进行修剪
m.num_features = bn_dict["backbone.layer1.0.bn1"].size()[0]

此外，还要对BN层前面的Conv的输出层进行修剪，对BN层后面的Conv的输入层进行修剪，整体代码：

#先得到所有BN层需要保留的权重索引
bn_dict = dict()
for name, m in model.named_modules():
    if isinstance(m, nn.BatchNorm2d):
        indice = find_indice(m, thresh=0.17)
        bn_dict[name] = indice
        
#进行剪枝
if name == "backbone.layer1.0.conv1":
    m.weight.data = m.weight.data[:, bn_dict["backbone.bn1"], :, :]
    m.weight.data = m.weight.data[bn_dict["backbone.layer1.0.bn1"], :, :, :]
    m.in_channels = bn_dict["backbone.bn1"].size()[0]
    m.out_channels = bn_dict["backbone.layer1.0.bn1"].size()[0]
if name == "backbone.layer1.0.bn1":
    m.weight.data = m.weight.data[bn_dict["backbone.layer1.0.bn1"]]  #gamma
    m.bias.data = m.bias.data[bn_dict["backbone.layer1.0.bn1"]]      #beta
    m.running_mean.data = m.running_mean.data[bn_dict["backbone.layer1.0.bn1"]]
    m.running_var.data = m.running_var.data[bn_dict["backbone.layer1.0.bn1"]]
    m.num_features = bn_dict["backbone.layer1.0.bn1"].size()[0]
if name == "backbone.layer1.0.conv2":
    m.weight.data = m.weight.data[:, bn_dict["backbone.layer1.0.bn1"], :, :]
    m.weight.data = m.weight.data[bn_dict["backbone.layer1.0.bn2"], :, :, :]
    m.in_channels = bn_dict["backbone.layer1.0.bn1"].size()[0]
    m.out_channels = bn_dict["backbone.layer1.0.bn2"].size()[0]

剪完之后再对模型进行保存，打印出来后的模型就是权重和结构都修剪好的模型：

torch.save(model, "修剪好的模型的保存路径")

打印：
在这里插入图片描述
这样修剪的模型也是可以直接运行的。

我自己的检测模型的backbone是resnet50，并且我只修剪了backbone部分，再剪完之后mAP下降了一些（和设定阈值有关，剪得多降得多），但是微调之后反而比剪之前的效果更好。我暂时把这个现象解释为，BN层的 $\gamma$ 参数可以当做一个Attention机制。如果有大佬能给出更好的解释，欢迎指正。

1.3注意事项

剪枝中间也才了一些坑，在此记录一下：

resnet50是有一些残差连接的，这是要注意一下的。就是在每一个reslayer层的最开始，又一个downsample层，由一个Conv和一个BN组成，如果该reslayer的输入的维度被剪了一部分，千万要记得对downsample层也进行调整！！！

附一张resnet50结构图：
在这里插入图片描述