简介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 层 gammaYOLOv5 通道重要性的天然打分器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_locationcpu) 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 的普通 ConvBN可剪占 FLOPs 大头剪枝收益最高C3 模块内部卷积分支可剪但需对齐分支经过 shortcut 拼接时两侧通道数必须一致SPPF 模块少剪内部池化不改变通道上下层卷积通道要同步重映射Detect 检测头通常不剪输出通道与类别数和 anchor 数量绑定剪完需重建后处理逻辑特别留意 C3 模块里的残差分支这是 YOLOv5 剪枝最容易掉进坑的地方。C3 中两条支路的输出会做 concatconcat 之后再接卷积如果只按 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_factor0.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_ratio0.3): prune_ratio0.3 表示剪掉全局 gamma 绝对值最小的 30% 通道。 返回 dictkey 为 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_tupleFalse).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_idxNone, out_idxNone): 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 下降mAP0.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_nameimages): 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_inputyolov5s.onnx, model_outputyolov5s_int8.onnx, calibration_data_readerYOLOv5CalibReader(calib_images), quant_formatQuantType.QInt8, per_channelTrue, )这段代码有三个坑要注意。第一CalibrationDataReader.get_next必须每次返回一个样本返回 None 表示结束不能用列表一次性全塞进去。第二input_name要和导出 ONNX 时实际输入名一致YOLOv5 默认导出输入名是images。第三per_channelTrue对卷积权重按输出通道分别算 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, inplaceFalse) # 此时可以直接复用 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%数据分布单一PTQ100~300 张校准图1% 以内剪枝率高30%类别多或场景变化大QAT全部训练集再跑一轮1%~2%部署后端为 TensorRTONNX 导出后 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能省下大量返工时间。本文还有配套的精品资源点击获取