思路：yolov11中的C3k2与yolov8的c2f的不同，所以与之前yolov8剪枝有稍许不同；
后续：会将剪枝流程写全，以及增加蒸馏、注意力、改loss；
注意：
1.在代码105行修改pruning.get_threshold(yolo.model, 0.65)，可以获得不同的剪枝率；
2.改代码放在训练代码同一页面下即可；
3.在最后修改文件夹地址来获得剪枝后的模型；
from ultralytics import YOLO
import torch
from ultralytics.nn.modules import Bottleneck, Conv, C2f, SPPF, Detect, C3k2
from torch.nn.modules.container import Sequential
import os# os.environ["CUDA_VISIBLE_DEVICES"] = "2"class PRUNE():def __init__(self) -> None:self.threshold = Nonedef get_threshold(self, model, factor=0.8):ws = []bs = []for name, m in model.named_modules():if isinstance(m, torch.nn.BatchNorm2d):w = m.weight.abs().detach()b = m.bias.abs().detach()ws.append(w)bs.append(b)print(name, w.max().item(), w.min().item(), b.max().item(), b.min().item())print()# keepws = torch.cat(ws)self.threshold = torch.sort(ws, descending=True)[0][int(len(ws) * factor)]def prune_conv(self, conv1: Conv, conv2: Conv):## a. 根据BN中的参数，获取需要保留的index================gamma = conv1.bn.weight.data.detach()beta = conv1.bn.bias.data.detach()keep_idxs = []local_threshold = self.thresholdwhile len(keep_idxs) < 8: ## 若剩余卷积核<8, 则降低阈值重新筛选keep_idxs = torch.where(gamma.abs() >= local_threshold)[0]local_threshold = local_threshold * 0.5n = len(keep_idxs)# n = max(int(len(idxs) * 0.8), p)print(n / len(gamma) * 100)# scale = len(idxs) / n## b. 利用index对BN进行剪枝============================conv1.bn.weight.data = gamma[keep_idxs]conv1.bn.bias.data = beta[keep_idxs]conv1.bn.running_var.data = conv1.bn.running_var.data[keep_idxs]conv1.bn.running_mean.data = conv1.bn.running_mean.data[keep_idxs]conv1.bn.num_features = nconv1.conv.weight.data = conv1.conv.weight.data[keep_idxs]conv1.conv.out_channels = nif isinstance(conv2, list) and len(conv2) > 3 and conv2[-1]._get_name() == "Proto":proto = conv2.pop()proto.cv1.conv.in_channels = nproto.cv1.conv.weight.data = proto.cv1.conv.weight.data[:, keep_idxs]## c. 利用index对conv1进行剪枝=========================if conv1.conv.bias is not None:conv1.conv.bias.data = conv1.conv.bias.data[keep_idxs]## d. 利用index对conv2进行剪枝=========================if not isinstance(conv2, list):conv2 = [conv2]for item in conv2:if item is None: continueif isinstance(item, Conv):conv = item.convelse:conv = itemif isinstance(item, Sequential):conv1 = item[0]conv = item[1].convconv1.conv.in_channels = nconv1.conv.out_channels = nconv1.conv.groups = nconv1.conv.weight.data = conv1.conv.weight.data[keep_idxs, :]conv1.bn.bias.data = conv1.bn.bias.data[keep_idxs]conv1.bn.weight.data = conv1.bn.weight.data[keep_idxs]conv1.bn.running_var.data = conv1.bn.running_var.data[keep_idxs]conv1.bn.running_mean.data = conv1.bn.running_mean.data[keep_idxs]conv1.bn.num_features = nconv.in_channels = nconv.weight.data = conv.weight.data[:, keep_idxs]def prune(self, m1, m2):if isinstance(m1, C3k2): # C2f as a top convm1 = m1.cv2if isinstance(m1, Sequential):m1 = m1[1]if not isinstance(m2, list): # m2 is just one modulem2 = [m2]for i, item in enumerate(m2):if isinstance(item, C3k2) or isinstance(item, SPPF):m2[i] = item.cv1self.prune_conv(m1, m2)def do_pruning(modelpath, savepath):pruning = PRUNE()### 0. 加载模型yolo = YOLO(modelpath) # build a new model from scratchpruning.get_threshold(yolo.model, 0.65) # 获取剪枝时bn参数的阈值，这里的0.8为剪枝率。### 1. 剪枝c2f 中的Bottleneckfor name, m in yolo.model.named_modules():if isinstance(m, Bottleneck):pruning.prune_conv(m.cv1, m.cv2)### 2. 指定剪枝不同模块之间的卷积核seq = yolo.model.modelfor i in [3, 5, 7, 8]:pruning.prune(seq[i], seq[i + 1])### 3. 对检测头进行剪枝# 在P3层: seq[15]之后的网络节点与其相连的有 seq[16]、detect.cv2[0] (box分支)、detect.cv3[0] (class分支)# 在P4层: seq[18]之后的网络节点与其相连的有 seq[19]、detect.cv2[1] 、detect.cv3[1]# 在P5层: seq[21]之后的网络节点与其相连的有 detect.cv2[2] 、detect.cv3[2]detect: Detect = seq[-1]proto = detect.protolast_inputs = [seq[16], seq[19], seq[22]]colasts = [seq[17], seq[20], None]for idx, (last_input, colast, cv2, cv3, cv4) in enumerate(zip(last_inputs, colasts, detect.cv2, detect.cv3, detect.cv4)):if idx == 0:pruning.prune(last_input, [colast, cv2[0], cv3[0], cv4[0], proto])else:pruning.prune(last_input, [colast, cv2[0], cv3[0], cv4[0]])pruning.prune(cv2[0], cv2[1])pruning.prune(cv2[1], cv2[2])pruning.prune(cv3[0], cv3[1])pruning.prune(cv3[1], cv3[2])pruning.prune(cv4[0], cv4[1])pruning.prune(cv4[1], cv4[2])### 4. 模型梯度设置与保存for name, p in yolo.model.named_parameters():p.requires_grad = Trueyolo.val(data='data.yaml', batch=2, device=0, workers=0)torch.save(yolo.ckpt, savepath)# yolo.model.pt_path = yolo.model.pt_path.replace("last.pt", os.path.basename(savepath))# yolo.export(format="onnx")## ## 重新load模型，修改保存命名，用以比较剪枝前后的onnx的大小# yolo = YOLO(modelpath) # build a new model from scratch# yolo.export(format="onnx")if __name__ == "__main__":modelpath = "runs/segment/Constraint/weights/best.pt"savepath = "runs/segment/Constraint/weights/last_prune.pt"do_pruning(modelpath, savepath)