📚 课程目录第4章 YOLOv5s模型剪枝与微调

4. YOLOv5s模型剪枝与微调

任务目标

对YOLOv5s模型执行通道级结构化剪枝（基于BN层γ稀疏性），移除贡献低的通道后微调，恢复因剪枝造成的精度损失，并对比剪枝前后模型文件大小与推理速度。

环境依赖

Python 3.8+，PyTorch ≥1.8，YOLOv5官方仓库（v6.0分支），thop库（用于FLOPs统计）。需提前准备COCO子集或自定义数据集（本案例使用VOC格式示例），并确保YOLOv5s预训练权重yolov5s.pt可访问。

实现方案

整体流程分为三步：稀疏化训练（在BN层γ上添加L1正则，使部分通道γ趋近于0）→ 结构化剪枝（按全局阈值裁剪低γ通道，生成窄网络）→ 微调恢复（加载剪枝后模型，用小学习率训练若干epoch）。核心代码如下：

# 文件: prune_yolov5s.py
import torch
import torch.nn as nn
import copy
from pathlib import Path
import sys

# 假设YOLOv5仓库位于 ../yolov5
sys.path.append('../yolov5')
from models.yolo import Model, Detect
from utils.torch_utils import model_info

# ---------- 1. 稀疏化训练（BN γ L1正则） ----------
def sparse_train(cfg='yolov5s.yaml', weights='yolov5s.pt', epochs=50, sr=0.001):
    model = Model(cfg, ch=3, nc=80).to('cuda')
    ckpt = torch.load(weights, map_location='cuda')
    model.load_state_dict(ckpt['model'].float().state_dict())
    
    # 在训练循环中，每步backward后对BN γ施加惩罚
    # 此处仅展示核心逻辑，实际需嵌入YOLOv5的train.py
    for name, m in model.named_modules():
        if isinstance(m, nn.BatchNorm2d):
            m.weight.grad.data.add_(sr * torch.sign(m.weight.data))  # L1梯度
    
    # 训练完成后保存稀疏化权重
    torch.save({'model': model}, 'sparse.pt')
    return model

# ---------- 2. 结构化剪枝 ----------
def prune_model(model, percent=0.3):
    model = copy.deepcopy(model)
    # 收集所有BN层γ值
    bn_weights = []
    for m in model.modules():
        if isinstance(m, nn.BatchNorm2d):
            bn_weights.extend(m.weight.detach().cpu().numpy().tolist())
    
    # 计算全局剪枝阈值
    bn_weights = torch.tensor(bn_weights)
    threshold = torch.quantile(bn_weights, percent)
    
    # 为每个BN层生成通道保留掩码
    prune_masks = {}
    for name, m in model.named_modules():
        if isinstance(m, nn.BatchNorm2d):
            mask = (m.weight.detach().abs() > threshold).float()
            prune_masks[name] = mask
    
    # 应用掩码：将γ置零（实际剪枝需重写卷积层，此处用掩码模拟）
    for name, m in model.named_modules():
        if isinstance(m, nn.BatchNorm2d):
            m.weight.data.mul_(prune_masks[name])
            m.bias.data.mul_(prune_masks[name])
    
    # 真实剪枝需调用 utils/prune.py 中的 prune_model 函数
    # 此处为演示，返回带掩码的模型
    return model, prune_masks

# ---------- 3. 微调 ----------
def finetune(model, epochs=30, lr=0.001):
    # 冻结除BN外的层，仅训练BN γ（或全量微调）
    for name, p in model.named_parameters():
        if 'bn' not in name:
            p.requires_grad = False
    
    optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=lr)
    # 训练循环（需接入数据加载器）
    # 此处仅展示核心逻辑
    return model

# ---------- 主流程 ----------
if __name__ == '__main__':
    # 步骤1：稀疏化训练（实际需运行完整训练脚本）
    # sparse_train()
    
    # 步骤2：加载稀疏化模型并剪枝
    ckpt = torch.load('sparse.pt', map_location='cuda')
    model = ckpt['model']
    pruned_model, masks = prune_model(model, percent=0.3)
    
    # 步骤3：微调（需接入数据）
    # finetune(pruned_model)
    
    # 比较模型大小
    torch.save(model.state_dict(), 'yolov5s_dense.pth')
    torch.save(pruned_model.state_dict(), 'yolov5s_pruned.pth')
    
    dense_size = Path('yolov5s_dense.pth').stat().st_size / 1e6
    pruned_size = Path('yolov5s_pruned.pth').stat().st_size / 1e6
    print(f'原始模型大小: {dense_size:.2f} MB')
    print(f'剪枝后模型大小: {pruned_size:.2f} MB')
    print(f'压缩率: {dense_size/pruned_size:.2f}x')

运行与验证

稀疏化训练：修改YOLOv5官方train.py，在loss.backward()后插入sparse_train中的γ梯度惩罚代码，运行python train.py --weights yolov5s.pt --epochs 50生成sparse.pt。
执行剪枝：运行python prune_yolov5s.py，观察输出的模型大小对比。
微调：将剪枝后模型接入train.py，设置--weights pruned.pt --epochs 30 --lr 0.001完成精度恢复。

通过条件：剪枝后模型mAP@0.5下降不超过2%（相对原始模型），模型文件体积减少≥30%，单张图片推理时间（GPU）缩短≥20%。

常见问题

剪枝后精度骤降：降低剪枝比例（如从30%调至20%），或增加微调epoch数。
γ值未趋近于0：增大稀疏化系数sr至0.005，或延长稀疏化训练epoch。
实际剪枝未生效：本案例演示了掩码逻辑，真实剪枝需调用YOLOv5仓库utils/prune.py中的prune_model函数，它会重写卷积层输入/输出通道数。