yolov11剪枝

思路：yolov11中的C3k2与yolov8的c2f的不同，以是与之前yolov8剪枝有稍许不同；

后续：会将剪枝流程写全，以及增加蒸馏、注意力、改loss；

注意：
1.在代码105行修改pruning.get_threshold(yolo.model, 0.65)，可以得到不同的剪枝率；
2.改代码放在训练代码同一页面下即可；
3.在最后修改文件夹地点来得到剪枝后的模型；

1. from ultralytics import YOLO
2. import torch
3. from ultralytics.nn.modules import Bottleneck, Conv, C2f, SPPF, Detect, C3k2
4. from torch.nn.modules.container import Sequential
5. import os
8. # os.environ["CUDA_VISIBLE_DEVICES"] = "2"
11. class PRUNE():
12.     def __init__(self) -> None:
13.         self.threshold = None
15.     def get_threshold(self, model, factor=0.8):
16.         ws = []
17.         bs = []
18.         for name, m in model.named_modules():
19.             if isinstance(m, torch.nn.BatchNorm2d):
20.                 w = m.weight.abs().detach()
21.                 b = m.bias.abs().detach()
22.                 ws.append(w)
23.                 bs.append(b)
24.                 print(name, w.max().item(), w.min().item(), b.max().item(), b.min().item())
25.                 print()
26.         # keep
27.         ws = torch.cat(ws)
28.         self.threshold = torch.sort(ws, descending=True)[0][int(len(ws) * factor)]
30.     def prune_conv(self, conv1: Conv, conv2: Conv):
31.         ## a. 根据BN中的参数，获取需要保留的index================
32.         gamma = conv1.bn.weight.data.detach()
33.         beta = conv1.bn.bias.data.detach()
35.         keep_idxs = []
36.         local_threshold = self.threshold
37.         while len(keep_idxs) < 8:  ## 若剩余卷积核<8, 则降低阈值重新筛选
38.             keep_idxs = torch.where(gamma.abs() >= local_threshold)[0]
39.             local_threshold = local_threshold * 0.5
40.         n = len(keep_idxs)
41.         # n = max(int(len(idxs) * 0.8), p)
42.         print(n / len(gamma) * 100)
43.         # scale = len(idxs) / n
45.         ## b. 利用index对BN进行剪枝============================
46.         conv1.bn.weight.data = gamma[keep_idxs]
47.         conv1.bn.bias.data = beta[keep_idxs]
48.         conv1.bn.running_var.data = conv1.bn.running_var.data[keep_idxs]
49.         conv1.bn.running_mean.data = conv1.bn.running_mean.data[keep_idxs]
50.         conv1.bn.num_features = n
51.         conv1.conv.weight.data = conv1.conv.weight.data[keep_idxs]
52.         conv1.conv.out_channels = n
54.         if isinstance(conv2, list) and len(conv2) > 3 and conv2[-1]._get_name() == "Proto":
55.             proto = conv2.pop()
56.             proto.cv1.conv.in_channels = n
57.             proto.cv1.conv.weight.data = proto.cv1.conv.weight.data[:, keep_idxs]
58.         ## c. 利用index对conv1进行剪枝=========================
59.         if conv1.conv.bias is not None:
60.             conv1.conv.bias.data = conv1.conv.bias.data[keep_idxs]
62.         ## d. 利用index对conv2进行剪枝=========================
63.         if not isinstance(conv2, list):
64.             conv2 = [conv2]
65.         for item in conv2:
66.             if item is None: continue
67.             if isinstance(item, Conv):
68.                 conv = item.conv
69.             else:
70.                 conv = item
71.             if isinstance(item, Sequential):
72.                 conv1 = item[0]
73.                 conv = item[1].conv
74.                 conv1.conv.in_channels = n
75.                 conv1.conv.out_channels = n
76.                 conv1.conv.groups = n
77.                 conv1.conv.weight.data = conv1.conv.weight.data[keep_idxs, :]
78.                 conv1.bn.bias.data = conv1.bn.bias.data[keep_idxs]
79.                 conv1.bn.weight.data = conv1.bn.weight.data[keep_idxs]
80.                 conv1.bn.running_var.data = conv1.bn.running_var.data[keep_idxs]
81.                 conv1.bn.running_mean.data = conv1.bn.running_mean.data[keep_idxs]
82.                 conv1.bn.num_features = n
83.             conv.in_channels = n
84.             conv.weight.data = conv.weight.data[:, keep_idxs]
86.     def prune(self, m1, m2):
87.         if isinstance(m1, C3k2):  # C2f as a top conv
88.             m1 = m1.cv2
89.         if isinstance(m1, Sequential):
90.             m1 = m1[1]
91.         if not isinstance(m2, list):  # m2 is just one module
92.             m2 = [m2]
93.         for i, item in enumerate(m2):
94.             if isinstance(item, C3k2) or isinstance(item, SPPF):
95.                 m2[i] = item.cv1
97.         self.prune_conv(m1, m2)
100. def do_pruning(modelpath, savepath):
101.     pruning = PRUNE()
103.     ### 0. 加载模型
104.     yolo = YOLO(modelpath)  # build a new model from scratch
105.     pruning.get_threshold(yolo.model, 0.65)  # 获取剪枝时bn参数的阈值，这里的0.8为剪枝率。
107.     ### 1. 剪枝c2f 中的Bottleneck
108.     for name, m in yolo.model.named_modules():
109.         if isinstance(m, Bottleneck):
110.             pruning.prune_conv(m.cv1, m.cv2)
112.     ### 2. 指定剪枝不同模块之间的卷积核
113.     seq = yolo.model.model
114.     for i in [3, 5, 7, 8]:
115.         pruning.prune(seq[i], seq[i + 1])
117.     ### 3. 对检测头进行剪枝
118.     # 在P3层: seq[15]之后的网络节点与其相连的有 seq[16]、detect.cv2[0] (box分支)、detect.cv3[0] (class分支)
119.     # 在P4层: seq[18]之后的网络节点与其相连的有 seq[19]、detect.cv2[1] 、detect.cv3[1]
120.     # 在P5层: seq[21]之后的网络节点与其相连的有 detect.cv2[2] 、detect.cv3[2]
121.     detect: Detect = seq[-1]
122.     proto = detect.proto
123.     last_inputs = [seq[16], seq[19], seq[22]]
124.     colasts = [seq[17], seq[20], None]
125.     for idx, (last_input, colast, cv2, cv3, cv4) in enumerate(zip(last_inputs, colasts, detect.cv2, detect.cv3, detect.cv4)):
126.         if idx == 0:
127.             pruning.prune(last_input, [colast, cv2[0], cv3[0], cv4[0], proto])
128.         else:
129.             pruning.prune(last_input, [colast, cv2[0], cv3[0], cv4[0]])
130.         pruning.prune(cv2[0], cv2[1])
131.         pruning.prune(cv2[1], cv2[2])
132.         pruning.prune(cv3[0], cv3[1])
133.         pruning.prune(cv3[1], cv3[2])
134.         pruning.prune(cv4[0], cv4[1])
135.         pruning.prune(cv4[1], cv4[2])
137.     ### 4. 模型梯度设置与保存
138.     for name, p in yolo.model.named_parameters():
139.         p.requires_grad = True
141.     yolo.val(data='data.yaml', batch=2, device=0, workers=0)
142.     torch.save(yolo.ckpt, savepath)
143.     # yolo.model.pt_path = yolo.model.pt_path.replace("last.pt", os.path.basename(savepath))
144.     # yolo.export(format="onnx")
145.     #
146.     # ## 重新load模型，修改保存命名，用以比较剪枝前后的onnx的大小
147.     # yolo = YOLO(modelpath)  # build a new model from scratch
148.     # yolo.export(format="onnx")
151. if __name__ == "__main__":
152.     modelpath = "runs/segment/Constraint/weights/best.pt"
153.     savepath = "runs/segment/Constraint/weights/last_prune.pt"
154.     do_pruning(modelpath, savepath)

复制代码

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。