自定义代码进行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 ← γ 1 x 1 ^ + β 1 ≡ B N γ 1 , β 1 ( x 1 ) y_{1} \leftarrow \gamma_1 \hat{x_{1}}+\beta_1 \equiv B N_{\gamma_1, \beta_1}\left(x_{1}\right) y1←γ1x1^+β1≡BNγ1,β1(x1)
所以参数 γ \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结构图: