YOLOv8剪枝全过程

1 约束训练

1.1 修改YOLOv8代码：

ultralytics/yolo/engine/trainer.py

添加内容：

# Backward
self.scaler.scale(self.loss).backward()

# ========== 新增 ==========
l1_lambda = 1e-2 * (1 - 0.9 * epoch / self.epochs)
for k, m in self.model.named_modules():
    if isinstance(m, nn.BatchNorm2d):
        m.weight.grad.data.add_(l1_lambda * torch.sign(m.weight.data))
        m.bias.grad.data.add_(1e-2 * torch.sign(m.bias.data))
# ========== 新增 ==========

# Optimize - https://pytorch.org/docs/master/notes/amp_examples.html
if ni - last_opt_step >= self.accumulate:
    self.optimizer_step()
    last_opt_step = ni

1.2  训练

需要注意的就是amp=False

命令行输入：
yolo train model=yolov8s.yaml epochs=100 amp=False

训练完会得到一个best.pt和last.pt，推荐用last.pt

1.3 约束训练可视化

已实现在tensorboard可视化约束训练过程BN参数的分布变化

随着训练进行（纵轴是epoch），BN层参数会逐渐从最上面的正太分布趋向于0附近。

以下是正常训练和稀疏训练的BN层参数值分布图：

右图的稀疏训练明显太早就全到0了，这样会影响精度，可以把系数1e-2改小一点1e-3，这样会稀疏的慢一点，l1_lambda = 1e-2 * (1 - 0.9 * epoch / self.epochs)
如下图：左为1e-2, 右为0.3*1e-2

2 剪枝

上一步得到的last.pt作为剪枝对象，自己创建一个prun.py文件：

这里的剪枝代码仅适用yolov8原模型，如有模块/模型的更改，则需要修改剪枝代码

需要定制改模型后的剪枝的可以私信

from ultralytics import YOLO
import torch
from ultralytics.nn.modules import Bottleneck, Conv, C2f, SPPF, Detect

# Load a model
yolo = YOLO("last.pt")
model = yolo.model

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())
# keep
factor = 0.8
ws = torch.cat(ws)
threshold = torch.sort(ws, descending=True)[0][int(len(ws) * factor)]
print(threshold)

def prune_conv(conv1: Conv, conv2: Conv):
    gamma = conv1.bn.weight.data.detach()
    beta = conv1.bn.bias.data.detach()
    keep_idxs = []
    local_threshold = threshold
    while len(keep_idxs) < 8:
        keep_idxs = torch.where(gamma.abs() >= local_threshold)[0]
        local_threshold = local_threshold * 0.5
    n = len(keep_idxs)
    # n = max(int(len(idxs) * 0.8), p)
    # print(n / len(gamma) * 100)
    # scale = len(idxs) / n
    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 = n
    conv1.conv.weight.data = conv1.conv.weight.data[keep_idxs]
    conv1.conv.out_channels = n

    if conv1.conv.bias is not None:
        conv1.conv.bias.data = conv1.conv.bias.data[keep_idxs]

    if not isinstance(conv2, list):
        conv2 = [conv2]

    for item in conv2:
        if item is not None:
            if isinstance(item, Conv):
                conv = item.conv
            else:
                conv = item
            conv.in_channels = n
            conv.weight.data = conv.weight.data[:, keep_idxs]


def prune(m1, m2):
    if isinstance(m1, C2f):  # C2f as a top conv
        m1 = m1.cv2

    if not isinstance(m2, list):  # m2 is just one module
        m2 = [m2]

    for i, item in enumerate(m2):
        if isinstance(item, C2f) or isinstance(item, SPPF):
            m2[i] = item.cv1

    prune_conv(m1, m2)


for name, m in model.named_modules():
    if isinstance(m, Bottleneck):
        prune_conv(m.cv1, m.cv2)

seq = model.model
for i in range(3, 9):
    if i in [6, 4, 9]: continue
    prune(seq[i], seq[i + 1])

detect: Detect = seq[-1]
last_inputs = [seq[15], seq[18], seq[21]]
colasts = [seq[16], seq[19], None]
for last_input, colast, cv2, cv3 in zip(last_inputs, colasts, detect.cv2, detect.cv3):
    prune(last_input, [colast, cv2[0], cv3[0]])
    prune(cv2[0], cv2[1])
    prune(cv2[1], cv2[2])
    prune(cv3[0], cv3[1])
    prune(cv3[1], cv3[2])

for name, p in yolo.model.named_parameters():
    p.requires_grad = True

# yolo.val() # 剪枝模型进行验证 yolo.val(workers=0)
yolo.export(format="onnx") # 导出为onnx文件
# yolo.train(data="VOC.yaml", epochs=100) # 剪枝后直接训练微调

torch.save(yolo.ckpt, "prune.pt")
print("done")

上述代码只需修改：

1. 最顶上的yolo = YOLO("last.pt")改为第一步约束训练得到的文件路径，一般为runs/detect/train/weights/last.pt

2. 最下面的torch.save(yolo.ckpt, "prune.pt")改为想要保存的路径

运行完会得到prune.pt和prune.onnx可以在netron.app网站拖入onnx文件查看是否剪枝成功了，成功的话可以看到某些通道数字为单数或者一些不规律的数字，如下图：

3 回调训练（finetune）

回调训练的唯一关键点就在于不让模型从yaml文件加载结构，直接加载pt文件

这里大佬给了两种方法，我尝试第一种方法无效，第二种成功了

3.1 首先要把第一步约束训练的代码注释掉

3.2  修改相关代码，使模型不加载yaml文件

修改位置：yolo/engine/trainer.py的443行左右

self.model = self.get_model(cfg=cfg, weights=weights, verbose=RANK == -1)  # calls Model(cfg, weights)

# ========== 新增该行代码 ==========
self.model = weights
# ========== 新增该行代码 ==========

return ckpt

3.3 修改完代码就可以进行finetun训练了

命令行输入：
yolo train model=prune.pt epochs=100

结果展示：

约束训练last.pt：

剪枝后的prune.pt：

回调后的finetune.pt：

可以看到精度损失很小，但是参数量和浮点运算量下去了很多，推理速度在cpu上测试是变快了的，gpu上好像没啥变化

剪枝前后各层通道数对比：