YOLOv11 模型压缩与性能优化实战指南
环境依赖与配置
在进行模型优化之前，需要确保运行环境满足以下核心依赖版本要求，以保证代码的兼容性与稳定性：
- Python 版本：3.10.16
- PyTorch 框架：2.6.0
- Ultralytics 库：8.3.28
- CUDA 版本：12.4
- 硬件建议：NVIDIA GeForce RTX 4090 或同等算力显卡
完整的依赖列表可通过项目配置文件获取，实际部署时仅需安装核心推理与训练包。
优化流程详解
本方案采用五阶段迭代策略，依次完成基准训练、稀疏约束、结构剪枝、微调恢复以及知识蒸馏。用户可根据实际需求通过注释控制执行阶段。
1. 基准模型训练 (Baseline Training)
首先加载预训练权重进行常规训练，确立性能基准。需配置预训练权重路径与数据集配置文件路径。
PROJECT_ROOT = os.getcwd()
BASE_CKPT_PATH = os.path.join(PROJECT_ROOT, "yolo11n.pt")
DATASET_CONFIG = os.path.join(PROJECT_ROOT, "data.yaml")
def run_baseline_training():
model = YOLO(BASE_CKPT_PATH)
model.train(data=DATASET_CONFIG, device="0", imgsz=640, epochs=50, batch=2)
2. 稀疏化约束训练 (Sparse Constraint Training)
为实现后续剪枝，需在 BatchNorm 层引入 L1 正则化约束，促使缩放因子稀疏化。此阶段需加载上一阶段的最佳权重，并修改训练器逻辑。
在训练循环中注入正则化逻辑，动态调整惩罚系数：
# 在训练器内部注入以下逻辑
l1_penalty = 1e-2 * (1 - 0.9 * current_epoch / total_epochs)
for module_name, module in model.named_modules():
if isinstance(module, nn.BatchNorm2d):
module.weight.grad.data.add_(l1_penalty * torch.sign(module.weight.data))
训练完成后，务必移除上述修改以避免影响后续步骤。
3. 结构剪枝 (Model Pruning)
根据稀疏化训练后的缩放因子，对模型通道进行剪除。设定剪枝比例参数，输入剪枝前的权重路径，生成剪枝后的模型文件。
SPARSITY_RATIO = 0.8
PRUNE_INPUT_PATH = os.path.join(CONSTRAINT_TRAIN_DIR, "weights/last.pt")
PRUNE_OUTPUT_PATH = os.path.join(CONSTRAINT_TRAIN_DIR, "weights/prune.pt")
def execute_pruning():
from utils.yolo.det_pruning import do_pruning
do_pruning(PRUNE_INPUT_PATH, PRUNE_OUTPUT_PATH, SPARSITY_RATIO)
4. 剪枝后微调 (Fine-tuning)
剪枝操作可能导致精度下降，需对剪枝后的模型进行恢复性训练。加载剪枝模型并冻结部分参数或全量微调。
FINETUNE_DIR = os.path.join(PROJECT_ROOT, "runs/detect/finetune")
def run_finetuning():
model = YOLO(PRUNE_OUTPUT_PATH)
for param in model.parameters():
param.requires_grad = True
model.train(data=DATASET_CONFIG, device="0", imgsz=640, epochs=200, batch=2, name=FINETUNE_DIR)
5. 知识蒸馏与注意力增强 (Distillation & Attention)
最后阶段结合知识蒸馏与注意力机制。将基准模型作为教师网络，微调后的剪枝模型作为学生网络，并在学生网络中嵌入注意力模块以提升特征提取能力。
def run_distillation():
target_layers = ["6", "8", "13", "16", "19", "22"]
teacher_net = YOLO(BASE_CKPT_PATH)
student_net = YOLO(os.path.join(FINETUNE_DIR, 'weights/best.pt'))
# 为学生模型添加注意力机制
student_net = add_attention(student_net)
student_net.train(
data=DATASET_CONFIG,
Distillation=teacher_net.model,
loss_type='mgd',
layers=target_layers,
amp=False,
imgsz=1280,
epochs=300,
batch=2,
device=0,
lr0=0.001,
name=os.path.join(PROJECT_ROOT, "runs/detect/student")
)
训练参数说明
- data: 数据集配置文件路径，包含类别与图像信息。
- device: 指定计算设备，例如 CUDA 设备 ID。
- imgsz: 输入图像的分辨率尺寸。
- epochs: 完整训练迭代的轮次。
- batch: 单次迭代处理的样本数量。
- Distillation: 指定教师模型实例，用于知识迁移。
- loss_type: 蒸馏损失函数类型，如 MGD (Masked Generative Distillation)。
- layers: 参与蒸馏损失计算的网络层索引。
完整执行脚本
from ultralytics import YOLO
import os
from utils.yolo.attention import add_attention
PROJECT_ROOT = os.getcwd()
BASE_CKPT_PATH = os.path.join(PROJECT_ROOT, "yolo11n.pt")
DATASET_CONFIG = os.path.join(PROJECT_ROOT, "data.yaml")
CONSTRAINT_TRAIN_DIR = os.path.join(PROJECT_ROOT, "runs/detect/Constraint")
BASELINE_CKPT = os.path.join(PROJECT_ROOT, 'runs/detect/train3/weights/best.pt')
SPARSITY_RATIO = 0.8
PRUNE_INPUT_PATH = os.path.join(CONSTRAINT_TRAIN_DIR, "weights/last.pt")
PRUNE_OUTPUT_PATH = os.path.join(CONSTRAINT_TRAIN_DIR, "weights/prune.pt")
FINETUNE_DIR = os.path.join(PROJECT_ROOT, "runs/detect/finetune")
STUDENT_OUTPUT_DIR = os.path.join(PROJECT_ROOT, "runs/detect/student")
def phase_initial_train():
model = YOLO(BASE_CKPT_PATH)
model.train(data=DATASET_CONFIG, device="0", imgsz=640, epochs=50, batch=2)
def phase_sparse_train():
model = YOLO(BASELINE_CKPT)
model.train(data=DATASET_CONFIG, device="0", imgsz=640, epochs=50, batch=2, amp=False, name=CONSTRAINT_TRAIN_DIR)
def phase_pruning():
from utils.yolo.det_pruning import do_pruning
do_pruning(PRUNE_INPUT_PATH, PRUNE_OUTPUT_PATH, SPARSITY_RATIO)
def phase_finetune():
model = YOLO(PRUNE_OUTPUT_PATH)
for param in model.parameters():
param.requires_grad = True
model.train(data=DATASET_CONFIG, device="0", imgsz=640, epochs=200, batch=2, name=FINETUNE_DIR)
def phase_distillation():
layers = ["6", "8", "13", "16", "19", "22"]
teacher_net = YOLO(BASELINE_CKPT)
student_net = YOLO(os.path.join(FINETUNE_DIR, 'weights/best.pt'))
student_net = add_attention(student_net)
student_net.train(data=DATASET_CONFIG, Distillation=teacher_net.model, loss_type='mgd', layers=layers,
amp=False, imgsz=1280, epochs=300, batch=2, device=0, lr0=0.001, name=STUDENT_OUTPUT_DIR)
if __name__ == '__main__':
# 按需取消注释以执行对应阶段
# phase_initial_train()
# phase_sparse_train()
# phase_pruning()
# phase_finetune()
phase_distillation()
