简介:YOLOv5剪枝与量化一键运行代码包,面向需要在移动端、嵌入式设备或NVIDIA GPU上高效部署目标检测模型的开发者与研究者,主要解决模型参数量大、推理速度慢、存储占用高的痛点。资源将模型压缩完整流程封装为可直接运行的脚本,覆盖结构剪枝、权重剪枝、量化感知训练与TensorRT格式转换,实测可将模型体积减少70%以上,并在精度与推理速度之间取得较好平衡。压缩包共208个文件,包含Python训练与推理脚本、YAML模型配置、C++与CUDA源码、ONNX与PyTorch权重、Dockerfile环境配置、Markdown说明文档以及测试图片、视频,整体约24.2MB,目录结构清晰,便于快速定位与修改。已有1883人学习使用,特别适合刚接触模型压缩的工程师通过现成代码快速上手。借助该资源,使用者无需深入底层细节,即可完成剪枝、量化并生成TensorRT引擎,缩短模型优化与上线周期;配置灵活的脚本和示例也有助于二次开发。
1. YOLOv5 部署占不满预算时,为什么剪枝和量化必须一起做
从 YOLOv5s 部署到边缘设备,过程比预想中要长不少。在 640×640 输入下,YOLOv5s 的 FP32 模型约 14 MB、16.5 GFLOPs,在 Jetson、RK3588 上能跑实时,放到低功耗 MCU 或共享推理服务上帧率立刻不可用。真正解决部署问题的关键,是把结构稀疏化和数值量化同时做:剪枝压通道数,量化把 FP32 换成 INT8,参数量、计算量、内存带宽一起变小,精度靠微调与校准稳住。这个标题的关键是“代码一键运行”:难点不在单个命令,而在把稀疏化训练、通道裁剪、微调、量化校准、精度验证整条链路串起来。下面按这条链路依次拆解,每一步都给出可直接套用的代码和参数。
2. YOLOv5 结构化剪枝原理:通道剪枝、BN 稀疏化与可剪层边界
2.1 通道剪枝为什么比非结构化剪枝更适合部署
模型剪枝分两类。非结构化剪枝把每个卷积核内部不重要的权重直接置零,精度损失小,但权重变成稀疏矩阵,推理时依赖专门的稀疏算子库才能提速;PyTorch 和 ONNX Runtime 对稀疏算子的覆盖不完整,GPU 上经常跑不出理想收益。结构化剪枝不同,它直接删掉整个输出通道,权重形状改变、模型体积变小,但算子类型不变,TensorRT、OpenVINO、ONNX Runtime 这些推理引擎不需要特殊算子就能接管。工程上要落地,首选必然是通道级剪枝。
通道级剪枝也不是没有代价。它要求剪掉的通道在后续所有层做索引重映射,YOLOv5 的 C3 模块和 concat 结构越多,重映射越容易出错。所以做剪枝之前,先把网络结构里哪些层可剪、哪些层牵一发动全身摸清楚,再写代码。
2.2 BN 层 gamma:YOLOv5 通道重要性的天然打分器
YOLOv5 网络结构中,绝大多数卷积层后面都跟着 BatchNorm。BatchNorm 的缩放系数 gamma 直接决定了这个通道的输出幅度:
y = gamma * (x - mean) / sqrt(var + eps) + beta
当某通道的 gamma 绝对值趋近于 0,无论输入怎么变化,该通道输出几乎不变,说明它对最终检测结果的贡献很小。把 gamma 绝对值作为通道重要性的依据,是当前对 YOLOv5 做剪枝最常用的做法,出自网络瘦身的思路,落地可靠。
做法上分两步:第一步,稀疏化训练,让一部分 gamma 向 0 收敛;第二步,根据 gamma 分布设全局阈值,把 gamma 绝对值小于阈值的通道删掉。先看第一步里如何判断稀疏化是否到位,可以用这段脚本观察训练后 BN 层 gamma 的分布:
# gamma_dist.py —— 统计 YOLOv5 权重中所有 BN 层 gamma 的绝对值分布 import torch ckpt = torch.load("runs/train/exp/weights/best.pt", map_location="cpu") model = ckpt["model"].float() gammas = [] for m in model.modules(): if isinstance(m, torch.nn.BatchNorm2d): gammas += m.weight.detach().abs().view(-1).tolist() gammas = sorted(gammas) print("gamma 数量:", len(gammas)) print("中位数:", gammas[len(gammas) // 2]) print("尾部 10%:", gammas[:int(len(gammas) * 0.1)]) print("最大值:", gammas[-1])ckpt["model"]对应 YOLOv5 保存权重时的键名,加载后要转回 float 再统计。如果尾部 10% 的 gamma 都接近 0,说明稀疏化训练起了作用,可以继续做剪枝;如果分布整体还是普通正态,说明稀疏化系数太小或训练 epoch 不足。
2.3 YOLOv5 网络结构中哪些层能剪、哪些不能碰
YOLOv5 由卷积块、C3 模块、SPPF 和 Detect 检测头组成。剪枝时我一般按这个表来定边界:
| 部位 | 是否可剪 | 原因 |
|---|---|---|
| 首个卷积层(stem) | 不剪或只剪少量 | 输入固定为 RGB 3 通道,剪掉后输入数据格式要跟着改,收益小 |
| 主干和 Neck 的普通 Conv+BN | 可剪 | 占 FLOPs 大头,剪枝收益最高 |
| C3 模块内部卷积分支 | 可剪但需对齐 | 分支经过 shortcut 拼接时,两侧通道数必须一致 |
| SPPF 模块 | 少剪 | 内部池化不改变通道,上下层卷积通道要同步重映射 |
| Detect 检测头 | 通常不剪 | 输出通道与类别数和 anchor 数量绑定,剪完需重建后处理逻辑 |
特别留意 C3 模块里的残差分支,这是 YOLOv5 剪枝最容易掉进坑的地方。C3 中两条支路的输出会做 concat,concat 之后再接卷积,如果只按 BN gamma 阈值分别剪两条支路,concat 后的通道拼接顺序和数量都要重新推算。常见做法是:从模型最前端开始,逐层维护一张「通道索引映射表」,每经过一个 concat 就把参与拼接的所有分支索引按原始顺序拼接,再传给下一层卷积做输入重映射。一旦顺序错位,模型推理出来就是乱码。
3. YOLOv5 剪枝落地代码:稀疏化训练、通道掩码重建与微调
3.1 稀疏化训练:在 YOLOv5 的 loss 上追加 L1 正则
标准的 YOLOv5 训练循环不会自动让 gamma 稀疏,需要在 loss 计算处插入一段稀疏正则代码。常见做法是直接改 train.py,在得到总 loss 之后、反向传播之前追加:
# train_sparse.py 核心片段 —— 在 YOLOv5 训练循环中追加稀疏损失 import torch import torch.nn as nn def add_sparse_loss(model, loss, s_factor=0.005): """ 对全部 BN 层的 gamma 施加 L1 正则。 s_factor 越大,gamma 稀疏化越激进,但过大会让精度崩掉。 """ sparse_loss = 0.0 for m in model.modules(): if isinstance(m, nn.BatchNorm2d): sparse_loss += torch.abs(m.weight).sum() return loss + s_factor * sparse_loss参数说明:s_factor常见区间是 0.001 到 0.01。COCO 这种大而全的数据集用 0.005 起步,自己的小数据集或类别数少时,可以提到 0.01。稀疏化训练的 epoch 也要拉长,建议是正常训练的 1.5 到 2 倍,让 gamma 的稀疏有个渐进过程。多卡训练时,这段损失在每张卡上各自算一遍,不需要额外做跨卡 reduce,反向传播的梯度同步会把稀疏正则的影响带进去。
训练完成后,用 2.2 节的统计脚本看 gamma 分布。一个常见误判是只看 loss 曲线下降,不看 gamma 分布,结果跑完根本不满足稀疏条件,后面剪枝等于盲剪。
3.2 通道掩码生成:全局阈值加索引收集
稀疏化训练完成之后,剪枝的核心就两步:算出哪些通道保留,再把保留索引传给模型重建逻辑。生成掩码的关键是「全局阈值」而不是逐层阈值,否则每层剪掉的比例不同,网络层的输出维度在层间传递时会乱套:
# prune.py —— 根据全局阈值生成每个 BN 层的保留通道索引 import numpy as np import torch import torch.nn as nn def build_channel_mask(model, prune_ratio=0.3): """ prune_ratio=0.3 表示剪掉全局 gamma 绝对值最小的 30% 通道。 返回 dict:key 为 BN 层在模型中的唯一名称,value 为保留通道索引列表。 """ alphas = [] for m in model.modules(): if isinstance(m, nn.BatchNorm2d): alphas.extend(m.weight.detach().abs().cpu().numpy()) threshold = np.percentile(alphas, prune_ratio * 100) masks = {} for name, m in model.named_modules(): if isinstance(m, nn.BatchNorm2d): keep = m.weight.detach().abs().cpu() > threshold masks[name] = keep.nonzero(as_tuple=False).view(-1).tolist() print(f"{name}: 保留 {len(masks[name])} / {len(keep)} 通道") return masks这段代码里threshold是一次性算出来的全局分位数。拿到masks之后还不能直接删层,必须按网络前向顺序逐层重排权重。真正做模型重建时,我的做法是维护一张全局索引表:每个卷积层的输入通道索引,由它的前一层输出 mask 决定;输出通道索引,由该卷积后跟的 BN 层 mask 决定;遇到 concat,就把所有分支的索引拼接起来再传给下一层。这个过程最容易出错,建议重建后立即用单张测试图对比原始模型和剪枝模型的输出特征图尺寸,提前发现维度不匹配。
3.3 模型重建与剪枝后微调
重建时,把保留下来的通道权重拷贝进新的模型结构。下面这段代码是重建模块里最核心的权重重排动作:
# rebuild.py 核心片段 —— 卷积权重的输入输出通道重排 def reindex_conv(conv, in_idx=None, out_idx=None): """ in_idx:输入通道保留索引,对应前一层 BN 的 mask out_idx:输出通道保留索引,对应本层后续 BN 的 mask """ w = conv.weight.data if out_idx is not None: w = w[out_idx] if in_idx is not None: w = w[:, in_idx] conv.weight.data = w if conv.bias is not None and out_idx is not None: conv.bias.data = conv.bias.data[out_idx]重建完成后,模型结构已经变了,要把新结构导出成独立的pruned.yaml,不能沿用原来的yolov5s.yaml。之后进入微调阶段,命令和普通训练类似,但参数要调整:
python train.py \ --data data/custom.yaml \ --cfg pruned.yaml \ --weights pruned.pt \ --epochs 80 \ --batch-size 16 \ --lr0 0.005 \ --freeze 10微调有四个要点:不再追加稀疏正则,让 gamma 自由恢复;冻结前 10 层,浅层语义已经稳定,只更新深层和检测头;学习率砍到正常训练的一半,比如从 0.01 降到 0.005;训练 epoch 不用太长,50 到 80 回合足够。剪枝比例和精度的关系可以按下面这个经验表预估:
| 剪枝比例 | FLOPs 下降 | mAP@0.5 相对基线 | 建议使用场景 |
|---|---|---|---|
| 20% | 约 20% | 98%~99% | 基本无损,适合直接上 |
| 30% | 约 30% | 95%~97% | 平衡点,多数场合常用 |
| 40% | 约 40% | 91%~94% | 对精度敏感的部署不建议 |
| 50% 以上 | 接近一半 | 85% 以下 | 仅能配合 QAT 和蒸馏抢救 |
注意:剪枝比例超过 40% 时,单靠微调很难把 mAP 拉回 95% 以上,除非后续配合第 4 章要讲的 QAT,再加上知识蒸馏才能兜住。
4. YOLOv5 量化落地代码:PTQ 校准、QAT 微调与精度验证
4.1 量化的本质:用 scale 和 zero_point 把 FP32 换 INT8
量化做的是数值重映射。对每个张量,找出一组浮点阈值范围 [min, max],用 scale 和 zero_point 把 FP32 数值映射成 INT8 整数。推理时卷积运算变成整数乘加,内存带宽和计算量同时下降。YOLOv5 的模型量化有两个层级:动态量化只量化权重,适合文本类模型;视觉模型要压推理延迟,必须静态量化,把激活值也量化到 INT8,这就需要准备一小批校准数据,逐个算子统计数值范围。
4.2 用 ONNX Runtime 做 PTQ 静态量化
YOLOv5 官方就有导出 ONNX 的能力,拿到 ONNX 模型后用 ONNX Runtime 的静态量化接口做 INT8 转换,是部署链路里最直接的做法。校准数据集建议 100 到 300 张,覆盖目标类别、不同尺度和不同光照,数量不是越多越好,而是要尽量贴近真实推理环境:
# quant_onnx.py —— 对导出后的 YOLOv5 ONNX 模型做 INT8 静态量化 import numpy as np from onnxruntime.quantization import ( quantize_static, QuantType, CalibrationDataReader ) class YOLOv5CalibReader(CalibrationDataReader): """按 ONNX Runtime 要求逐个喂入校准图片""" def __init__(self, calib_images, input_name="images"): self.data = [{ input_name: img.astype(np.float32) } for img in calib_images] self.iter_id = 0 def get_next(self): if self.iter_id < len(self.data): sample = self.data[self.iter_id] self.iter_id += 1 return sample return None quantize_static( model_input="yolov5s.onnx", model_output="yolov5s_int8.onnx", calibration_data_reader=YOLOv5CalibReader(calib_images), quant_format=QuantType.QInt8, per_channel=True, )这段代码有三个坑要注意。第一,CalibrationDataReader.get_next必须每次返回一个样本,返回 None 表示结束,不能用列表一次性全塞进去。第二,input_name要和导出 ONNX 时实际输入名一致,YOLOv5 默认导出输入名是images。第三,per_channel=True对卷积权重按输出通道分别算 scale,精度通常比 per-tensor 高,但对部分推理引擎指令集有要求,如果目标设备不支持,要先跑通 per-tensor 再切换。
ONNX 导出时的 opset 版本也要对。opset 12 以下对量化支持不完整,部分 Conv 和 MatMul 算子会被跳过,导致量化后模型甚至比原模型还大。
PTQ 完成之后,比较 INT8 模型和 FP32 模型在同一批验证集上的 mAP。剪枝率较低时,PTQ 的精度损失通常能控制在 1% 以内。如果损失超过 2%,就要考虑 4.3 节的 QAT 方案。
4.3 剪枝加量化掉精度,用 QAT 兜底
剪枝已经动过一次结构,再叠加 PTQ 的数值误差,容易出现精度跌破红线的情况。这时要上量化感知训练,让模型在训练阶段就把 INT8 的舍入误差学进去:
# qat.py —— 对剪枝后模型插入伪量化节点 import torch from torch.ao.quantization import ( prepare_qat, convert, get_default_qat_qconfig ) model.qconfig = get_default_qat_qconfig("fbgemm") model.train() # QAT 必须保持训练模式,BN 统计量继续更新 model_prepared = prepare_qat(model, inplace=False) # 此时可以直接复用 YOLOv5 的 train.py 流程,跑一轮正常训练 # 训练结束后执行: model_int8 = convert(model_prepared) torch.save(model_int8.state_dict(), "yolov5s_qat.pth")注意 QAT 训练不能在微调后的模型上继续加稀疏正则,两者会互相放大数值误差,最好先保存一份干净微调权重,再单独走 QAT。另外,QAT 之后还需要重新导出 ONNX,如果中途有版本不兼容的伪量化算子,导出会直接报错;先小步验证单个算子,再全模型导出。
PTQ 和 QAT 的选型对部署稳定性影响很大,遇到精度不达标时按下面这张表快速决策:
| 场景 | 推荐方式 | 校准/训练数据 | 预期精度损失 |
|---|---|---|---|
| 剪枝率低(10%~20%),数据分布单一 | PTQ | 100~300 张校准图 | 1% 以内 |
| 剪枝率高(30%),类别多或场景变化大 | QAT | 全部训练集再跑一轮 | 1%~2% |
| 部署后端为 TensorRT | ONNX 导出后 trtexec 做 INT8 | 校准图 | 与 PTQ/QAT 相当 |
4.4 量化后的两层验证
量化完不要只看 mAP。第一层验证是数值一致性:拿同一张图分别跑 FP32 和 INT8 模型,比较输出张量的余弦相似度,低于 0.95 就说明某层量化范围设置有问题,回到校准集检查是否覆盖不足。第二层验证是端到端指标:验证集完整跑一遍 mAP,和微调后模型做对比,相对下降超过 1% 就要评估是接受还是回退。
5. 一键运行编排:YAML 配置、Shell 串行编排与部署后三个检查
5.1 用 YAML 把整套流程参数统一收口
一键运行的真实痛点不是少执行一条命令,而是每次复制命令时漏掉参数。我的做法是把稀疏化系数、剪枝比例、微调 epoch、量化方式和校准图片路径全部收进一个配置文件:
# config.yaml data: data_yaml: ./data/custom.yaml model_cfg: yolov5s.yaml weights: ./weights/yolov5s.pt sparse: s_factor: 0.005 epochs: 150 prune: ratio: 0.3 finetune: epochs: 80 lr0: 0.005 freeze_layers: 10 quantize: method: ptq calib_images: 200 per_channel: true这套配置的好处是每个阶段独立成节,调参只动 YAML,代码逻辑不动。剪枝失败时也能根据 YAML 中的参数快速复现问题。
5.2 Shell 入口脚本:串起四个阶段的执行顺序
入口脚本要做的不是把命令堆叠,而是用set -euo pipefail保证任何一步失败就立刻中断,别让后续步骤拿着坏产物继续跑:
#!/usr/bin/env bash # run_all.sh —— YOLOv5 剪枝量化一键运行 set -euo pipefail echo "[1/4] 稀疏化训练" python train_sparse.py --config config.yaml echo "[2/4] 通道剪枝与模型重建" python rebuild.py --config config.yaml echo "[3/4] 剪枝后微调" python train.py --config config.yaml --phase finetune echo "[4/4] 量化与精度校验" python quantize_onnx.py --config config.yaml python eval_quantized.py --config config.yamlset -euo pipefail里特别注意pipefail要把第一个失败的管道命令状态作为返回值,避免后续的awk或tee把错误吞掉。每一步执行后检查产物文件是否生成,比如pruned.pt和yolov5s_int8.onnx,没有就直接退出。
5.3 部署后必须检查的三类情况
第一个检查:重建后的模型跑一次推理看维度。剪枝最容易出现的问题是 concat 对齐错位,拿一张测试图比较剪枝前后特征图尺寸,输出 shape 对不上就回头检查索引映射表。
第二个检查:量化后输出张量相似度。校准集如果只覆盖了白天样本,夜间推理精度必然下降,把校准集扩到目标场景的 2 到 3 倍再试。
第三个检查:中间产物一致性。pruned.yaml导出的结构和pruned.pt里的权重必须匹配,否则加载时报错或通道错位;每次剪枝后都重新生成 yaml,不要沿用旧文件。
最后建议在验证集里固定抽 20 张有代表性的图,量化后跑一次余弦相似度回归,把这个回归用例放进 CI。链路每次改动之后,先看这 20 张图是否过线,再看完整 mAP,能省下大量返工时间。
本文还有配套的精品资源,点击获取