📚 课程目录第1章 YOLOv11模型剪枝量化压缩实战100例

1. YOLOv11模型剪枝量化压缩实战100例

任务目标

本案例以YOLOv11检测模型为对象，完成从训练后权重到可部署轻量模型的完整压缩闭环：首先进行结构化剪枝（移除对检测贡献低的通道），随后进行INT8量化（将FP32权重压缩为INT8以减小体积并加速推理）。最终交付压缩后的模型文件、体积对比数据以及精度验证脚本，使模型参数量减少且推理速度提升，同时mAP@0.5下降不超过2%。

环境与依赖

# Python 3.9+，需安装以下核心库
pip install ultralytics==8.3.0 torch==2.1.0 onnx==1.16.0 onnxruntime==1.18.0

实现思路与关键代码

整体流程分为三步：① 通道稀疏化训练（为剪枝做准备）→ ② 结构化剪枝与微调 → ③ INT8量化导出。以下代码在单个脚本中完整实现该流程，输入为COCO格式训练数据，输出为压缩后的ONNX模型。

import torch
import torch.nn as nn
from ultralytics import YOLO
from ultralytics.nn.tasks import DetectionModel
import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType

# ---------- 1. 通道稀疏化训练（为剪枝制造通道重要性差异） ----------
model = YOLO('yolo11n.pt')  # 使用官方预训练权重
model.train(data='coco8.yaml', epochs=5, imgsz=640, 
            lr0=0.001, optimizer='SGD', 
            sparsity=0.001,  # 关键：启用L1通道稀疏正则
            sparsity_decay=0.0001)

# ---------- 2. 结构化剪枝：按BN层gamma值裁剪通道 ----------
def prune_channels(model, prune_ratio=0.3):
    """按BN层gamma值排序，裁剪最低30%通道"""
    bn_modules = []
    for m in model.model.modules():
        if isinstance(m, nn.BatchNorm2d):
            bn_modules.append(m)
    
    # 收集所有gamma值并计算全局阈值
    all_gamma = torch.cat([bn.weight.data.view(-1) for bn in bn_modules])
    threshold = torch.quantile(all_gamma.abs(), prune_ratio)
    
    # 生成剪枝掩码
    masks = {}
    for name, bn in zip([f'bn_{i}' for i in range(len(bn_modules))], bn_modules):
        mask = bn.weight.data.abs() > threshold
        masks[name] = mask
    
    # 应用剪枝（简化版：直接置零低重要性通道，实际工程需物理移除）
    for bn, mask in zip(bn_modules, masks.values()):
        bn.weight.data[~mask] = 0
        bn.bias.data[~mask] = 0
    
    return model

# 执行剪枝并微调恢复精度
pruned_model = prune_channels(model.model, prune_ratio=0.3)
model.model = pruned_model
model.train(data='coco8.yaml', epochs=3, imgsz=640, lr0=0.0005)  # 微调

# ---------- 3. 导出ONNX并INT8量化 ----------
model.export(format='onnx', imgsz=640, dynamic=False)  # 导出FP32 ONNX

# 动态INT8量化（无需校准数据，适合快速验证）
quantize_dynamic(
    model='yolo11n.onnx',
    output='yolo11n_int8.onnx',
    weight_type=QuantType.QInt8
)

# ---------- 4. 验证：对比压缩前后模型体积与精度 ----------
import os
fp32_size = os.path.getsize('yolo11n.onnx') / 1024 / 1024
int8_size = os.path.getsize('yolo11n_int8.onnx') / 1024 / 1024
print(f'FP32模型体积: {fp32_size:.2f} MB')
print(f'INT8模型体积: {int8_size:.2f} MB')
print(f'压缩比: {fp32_size / int8_size:.2f}x')

# 精度验证（使用验证集）
from ultralytics import YOLO as YOLOVal
val_model = YOLO('yolo11n_int8.onnx')
metrics = val_model.val(data='coco8.yaml', imgsz=640)
print(f'压缩后mAP@0.5: {metrics.box.map50:.4f}')

运行与验证方法

运行方式：将上述代码保存为compress_yolo11.py，在项目根目录执行python compress_yolo11.py。脚本会自动下载yolo11n.pt权重和coco8.yaml数据集（约1MB，用于快速验证流程）。

通过条件：满足以下全部条件即判定压缩成功：

脚本正常结束无报错，输出INT8模型文件yolo11n_int8.onnx；
INT8模型体积较FP32减小约3~4倍（动态量化典型压缩比）；
压缩后模型在coco8验证集上mAP@0.5不低于原始模型的98%（即下降≤2%）。

常见问题

剪枝后精度骤降：降低prune_ratio至0.2，并增加微调epochs至5~10；
量化后推理报错：确认onnxruntime版本≥1.18，且输入张量维度为固定640×640（不支持动态尺寸）；
稀疏训练不生效：检查ultralytics版本是否为8.3.0，且训练参数中sparsity值大于0。