yolov5模型压缩之模型剪枝_yolo剪枝

目前看来，yolo系列是工程上使用最为广泛的检测模型之一。yolov5检测性能优秀，部署便捷，备受广大开发者好评。但是，当模型在前端运行时，对模型尺寸与推理时间要求苛刻，轻量型模型yolov5s也难以招架。为了提高模型效率，这里与大家分享基于yolov5的模型剪枝方法 github分享连接。

剪枝原理与pipeline

本次使用稀疏训练对channel维度进行剪枝，来自论文Learning Efficient Convolutional Networks Through Network Slimming。其实原理很容易理解，我们知道bn层中存在两个可训练参数$\gamma ,\beta$，输入经过bn获得归一化后的分布。当$\gamma ，\beta$趋于0时，输入相当于乘上了0，那么，该channel上的卷积将只能输出0，毫无意义。因此，我们可以认为剔除这样的冗余channel对模型性能影响甚微。普通网络训练时，由于初始化，$\gamma$一般分布在1附近。为了使$\gamma$趋于0，可以通过添加L1正则来约束，使得系数稀疏化，论文中将添加$\gamma$L1正则的训练称为稀疏训练。
整个剪枝的过程如下图所示，首先初始化网络，对bn层的参数添加L1正则并对网络训练。统计网络中的$\gamma$，设置剪枝率对网络进行裁剪。最后，将裁减完的网络finetune，完成剪枝工作。

剪枝细节讲解

1.稀疏训练
上一章介绍了稀疏训练的原理，下面看一下代码是如何实现的。代码如下所示，首先，我们需要设置稀疏系数，稀疏系数对整个网络剪枝性能至关重要，设置太小的系数，$\gamma$趋于0的程度不高，无法对网络进行高强度的剪枝，但设置过大，会影响网络性能，大幅降低map。因此，我们需要通过实验找到合适的稀疏系数。
bn层的训练参数包括$\gamma ,\beta$，即代码中的m.weight,m.bias，loss.backward之后，在这两个参数的梯度上添加L1正则的梯度即可。

srtmp = opt.sr * (1 - 0.9 * epoch/epochs)
for k, m in model.named_modules():             
    if isinstance(m, nn.BatchNorm2d) and (k not in ignore_bn_list):
         m.weight.grad.data.add_(srtmp * torch.sign(m.weight.data))  # L1
         m.bias.grad.data.add_(opt.sr*10 * torch.sign(m.bias.data))  # L1
1
2
3
4
5

2.网络裁剪
上一步获得稀疏训练后的网络，接下来，我们需要将$\gamma$趋于0的channel裁剪掉。首先，统计所有BN层的$\gamma$，并对齐排序，找到剪枝率对应的阈值thre。

for i, layer in model.named_modules():
        if isinstance(layer, nn.BatchNorm2d):
            if i not in ignore_bn_list:
                model_list[i] = layer
            # bnw = layer.state_dict()['weight']
    model_list = {k:v for k,v in model_list.items() if k not in ignore_bn_list}
    prune_conv_list = [layer.replace("bn", "conv") for layer in model_list.keys()]
    bn_weights = gather_bn_weights(model_list)
    sorted_bn = torch.sort(bn_weights)[0]
    thre_index = int(len(sorted_bn) * opt.percent)
    thre = sorted_bn[thre_index]
1
2
3
4
5
6
7
8
9
10
11

然后，根据阈值获取每一bn层的mask，这里加了一些逻辑，目的是让剪枝后的channel保证是4的倍数，即复合前端加速要求。

def obtain_bn_mask(bn_module, thre):

    thre = thre.cuda()
    bn_layer = bn_module.weight.data.abs()
    temp = abs(torch.sort(bn_layer)[0][3::4] - thre)
    thre_temp = torch.sort(bn_layer)[0][3::4][temp.argmin()] 
    if int(temp.argmin()) == 0 and thre_temp > thre:
        thre = -1
    else:
        thre = thre_temp
    thre_index = int(bn_layer.shape[0] * 0.9)
    if thre_index % 4 != 0:
        thre_index -= thre_index % 4
    thre_perbn = torch.sort(bn_layer)[0][thre_index - 1]
    if thre_perbn < thre:
        thre = min(thre, thre_perbn)
    mask = bn_module.weight.data.abs().gt(thre).float()

    return mask
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

由于，剪枝后的网络与原网络channel不能对齐，因此，我们需要重新定义网络，并解析网络。重构的网络结构需要重新定义，因为需要导入更多的参数。

pruned_yaml["backbone"] =[[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
        [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
        [-1, 3, C3Pruned, [128]],
        [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
        [-1, 6, C3Pruned, [256]],
        [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
        [-1, 9, C3Pruned, [512]],
        [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
        [-1, 3, C3Pruned, [1024]],
        [-1, 1, SPPFPruned, [1024, 5]],  # 9
        ]
    pruned_yaml["head"] = [
        [-1, 1, Conv, [512, 1, 1]],
        [-1, 1, nn.Upsample, [None, 2, 'nearest']],
        [[-1, 6], 1, Concat, [1]],  # cat backbone P4
        [-1, 3, C3Pruned, [512, False]],  # 13

        [-1, 1, Conv, [256, 1, 1]],
        [-1, 1, nn.Upsample, [None, 2, 'nearest']],
        [[-1, 4], 1, Concat, [1]],  # cat backbone P3
        [-1, 3, C3Pruned, [256, False]],  # 17 (P3/8-small)

        [-1, 1, Conv, [256, 3, 2]],
        [[-1, 14], 1, Concat, [1]],  # cat head P4
        [-1, 3, C3Pruned, [512, False]],  # 20 (P4/16-medium)

        [-1, 1, Conv, [512, 3, 2]],
        [[-1, 10], 1, Concat, [1]],  # cat head P5
        [-1, 3, C3Pruned, [1024, False]],  # 23 (P5/32-large)

        [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
    ]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

yolov5的backbone与neck存在C3结构，C3中存在shortcut，即存在两个卷积相加的形式。为了使网络能够正常add，我们需要对add的两个卷积mask进行merge操作。与此同时，网络存在concate，所以还需要记录concate来自于哪些层以及concate输出的层。

for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']):  # from, number, module, args
        m = eval(m) if isinstance(m, str) else m  # eval strings
        for j, a in enumerate(args):
            try:
                args[j] = eval(a) if isinstance(a, str) else a  # eval strings
            except NameError:
                pass

        n = n_ = max(round(n * gd), 1) if n > 1 else n  # depth gain
        named_m_base = "model.{}".format(i)
        if m in [Conv]:
            named_m_bn = named_m_base + ".bn"

            bnc = int(maskbndict[named_m_bn].sum())
            c1, c2 = ch[f], bnc
            args = [c1, c2, *args[1:]]
            layertmp = named_m_bn
            if i>0:
                from_to_map[layertmp] = fromlayer[f]
            fromlayer.append(named_m_bn)

        elif m in [C3Pruned]:
            named_m_cv1_bn = named_m_base + ".cv1.bn"
            named_m_cv2_bn = named_m_base + ".cv2.bn"
            named_m_cv3_bn = named_m_base + ".cv3.bn"
            from_to_map[named_m_cv1_bn] = fromlayer[f]
            from_to_map[named_m_cv2_bn] = fromlayer[f]
            fromlayer.append(named_m_cv3_bn)
            
            if len(args) == 1:
                temp_mask = maskbndict[named_m_cv1_bn].bool() | maskbndict[named_m_base + '.m.0.cv2.bn'].bool()
                maskbndict[named_m_cv1_bn], maskbndict[named_m_base + '.m.0.cv2.bn'] = temp_mask.float(), temp_mask.float()
                
                if n > 1:
                    for repeat_ind in range(1, n):
                        temp_mask |= maskbndict[named_m_base + ".m.{}.cv2.bn".format(repeat_ind)].bool() 
                    for re_ind in range(n):
                        maskbndict[named_m_base + ".m.{}.cv2.bn".format(re_ind)] = temp_mask
                    maskbndict[named_m_cv1_bn], maskbndict[named_m_base + '.m.0.cv2.bn'] = temp_mask.float(), temp_mask.float()

            cv1in = ch[f]
            cv1out = int(maskbndict[named_m_cv1_bn].sum())
            cv2out = int(maskbndict[named_m_cv2_bn].sum())
            cv3out = int(maskbndict[named_m_cv3_bn].sum())
            args = [cv1in, cv1out, cv2out, cv3out, n, args[-1]]
            bottle_args = []
            chin = [cv1out]

            c3fromlayer = [named_m_cv1_bn]
            
            for p in range(n):
                named_m_bottle_cv1_bn = named_m_base + ".m.{}.cv1.bn".format(p)
                named_m_bottle_cv2_bn = named_m_base + ".m.{}.cv2.bn".format(p)
                bottle_cv1in = chin[-1]
                bottle_cv1out = int(maskbndict[named_m_bottle_cv1_bn].sum())
                bottle_cv2out = int(maskbndict[named_m_bottle_cv2_bn].sum())
                chin.append(bottle_cv2out)
                bottle_args.append([bottle_cv1in, bottle_cv1out, bottle_cv2out])
                from_to_map[named_m_bottle_cv1_bn] = c3fromlayer[p]
                from_to_map[named_m_bottle_cv2_bn] = named_m_bottle_cv1_bn
                c3fromlayer.append(named_m_bottle_cv2_bn)
            args.insert(4, bottle_args)
            c2 = cv3out
            n = 1
            from_to_map[named_m_cv3_bn] = [c3fromlayer[-1], named_m_cv2_bn]
        elif m in [SPPFPruned]:
            named_m_cv1_bn = named_m_base + ".cv1.bn"
            named_m_cv2_bn = named_m_base + ".cv2.bn"
            cv1in = ch[f]
            from_to_map[named_m_cv1_bn] = fromlayer[f]
            from_to_map[named_m_cv2_bn] = [named_m_cv1_bn]*4
            fromlayer.append(named_m_cv2_bn)
            cv1out = int(maskbndict[named_m_cv1_bn].sum())
            cv2out = int(maskbndict[named_m_cv2_bn].sum())
            args = [cv1in, cv1out, cv2out, *args[1:]]
            c2 = cv2out

        elif m is nn.BatchNorm2d:
            args = [ch[f]]
        elif m is Concat:
            c2 = sum(ch[x] for x in f)
            inputtmp = [fromlayer[x] for x in f]
            fromlayer.append(inputtmp)
        elif m is Detect:
            from_to_map[named_m_base + ".m.0"] = fromlayer[f[0]]
            from_to_map[named_m_base + ".m.1"] = fromlayer[f[1]]
            from_to_map[named_m_base + ".m.2"] = fromlayer[f[2]]
            args.append([ch[x] for x in f])
            if isinstance(args[1], int):  # number of anchors
                args[1] = [list(range(args[1] * 2))] * len(f)
        elif m is Contract:
            c2 = ch[f] * args[0] ** 2
        elif m is Expand:
            c2 = ch[f] // args[0] ** 2
        else:
            c2 = ch[f]
            fromtmp = fromlayer[-1]
            fromlayer.append(fromtmp)

        m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)  # module
        t = str(m)[8:-2].replace('__main__.', '')  # module type
        np = sum(x.numel() for x in m_.parameters())  # number params
        m_.i, m_.f, m_.type, m_.np = i, f, t, np  # attach index, 'from' index, type, number params
        save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelist
        layers.append(m_)
        if i == 0:
            ch = []
        ch.append(c2)
    return nn.Sequential(*layers), sorted(save), from_to_map
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109

重构并解析网络后，我们需要对解析后的网络填充参数，即找到解析后网络对应于原网络的各层参数，并clone赋值给重构后的网络，代码如下：

for ((layername, layer),(pruned_layername, pruned_layer)) in zip(model.named_modules(), pruned_model.named_modules()):
        assert layername == pruned_layername
        if isinstance(layer, nn.Conv2d) and not layername.startswith("model.24"):
            convname = layername[:-4]+"bn"
            if convname in from_to_map.keys():
                former = from_to_map[convname]
                if isinstance(former, str):
                    out_idx = np.squeeze(np.argwhere(np.asarray(maskbndict[layername[:-4] + "bn"].cpu().numpy())))
                    in_idx = np.squeeze(np.argwhere(np.asarray(maskbndict[former].cpu().numpy())))
                    w = layer.weight.data[:, in_idx, :, :].clone()
                    
                    if len(w.shape) ==3:     # remain only 1 channel.
                        w = w.unsqueeze(1)
                    w = w[out_idx, :, :, :].clone()
                    
                    pruned_layer.weight.data = w.clone()
                    changed_state.append(layername + ".weight")
                if isinstance(former, list):
                    orignin = [modelstate[i+".weight"].shape[0] for i in former]
                    formerin = []
                    for it in range(len(former)):
                        name = former[it]
                        tmp = [i for i in range(maskbndict[name].shape[0]) if maskbndict[name][i] == 1]
                        if it > 0:
                            tmp = [k + sum(orignin[:it]) for k in tmp]
                        formerin.extend(tmp)
                    out_idx = np.squeeze(np.argwhere(np.asarray(maskbndict[layername[:-4] + "bn"].cpu().numpy())))
                    w = layer.weight.data[out_idx, :, :, :].clone()
                    pruned_layer.weight.data = w[:,formerin, :, :].clone()
                    changed_state.append(layername + ".weight")
            else:
                out_idx = np.squeeze(np.argwhere(np.asarray(maskbndict[layername[:-4] + "bn"].cpu().numpy())))
                w = layer.weight.data[out_idx, :, :, :].clone()
                assert len(w.shape) == 4
                pruned_layer.weight.data = w.clone()
                changed_state.append(layername + ".weight")

        if isinstance(layer,nn.BatchNorm2d):
            out_idx = np.squeeze(np.argwhere(np.asarray(maskbndict[layername].cpu().numpy())))
            pruned_layer.weight.data = layer.weight.data[out_idx].clone()
            pruned_layer.bias.data = layer.bias.data[out_idx].clone()
            pruned_layer.running_mean = layer.running_mean[out_idx].clone()
            pruned_layer.running_var = layer.running_var[out_idx].clone()
            changed_state.append(layername + ".weight")
            changed_state.append(layername + ".bias")
            changed_state.append(layername + ".running_mean")
            changed_state.append(layername + ".running_var")
            changed_state.append(layername + ".num_batches_tracked")

        if isinstance(layer, nn.Conv2d) and layername.startswith("model.24"):
            former = from_to_map[layername]
            in_idx = np.squeeze(np.argwhere(np.asarray(maskbndict[former].cpu().numpy())))
            pruned_layer.weight.data = layer.weight.data[:, in_idx, :, :]
            pruned_layer.bias.data = layer.bias.data
            changed_state.append(layername + ".weight")
            changed_state.append(layername + ".bias")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

至此，我们完成了剪枝的所有步骤。