简介面向目标检测开发者的YOLOv5模型压缩实操包围绕剪枝、量化与TensorRT部署展开解决模型在移动端或资源受限设备上体积大、推理慢的问题。压缩包共208个文件约24.2MB以Python脚本、YAML配置、C/CUDA源码、Shell脚本和Dockerfile等为主同时包含ONNX、WTS格式的模型文件与说明文档便于按模块理解剪枝比例设置、量化感知训练和推理转换流程。已有1885人学习下载。整套代码提供一键运行脚本无需深入底层细节即可复现“剪枝减少70%以上参数、量化并转为TensorRT引擎”的完整链路。代码附带可配置超参数文件可调整剪枝力度也适合作为模型加速改造的参考模板直接迁移到实际部署项目。剪枝删除冗余连接和通道量化采用感知训练维持精度两者结合可大幅压缩模型压缩包内还提供示例图片与权重文件方便对比效果。1. 这个 YOLOv5 剪枝量化包解决的是你手里的哪类部署难题YOLOv5 训练出来的模型好跑但往 ncnn、RKNN、TensorRT 这些推理后端迁移的时候模型体积和延迟往往才是真正的坎。这个压缩包把剪枝、量化感知训练和 TensorRT 部署串成了一条龙剪枝比例能压到 70% 以上配合 INT8 量化单张显卡上能做到实时推理。代码不只是训练端的 Python 脚本连 TensorRT 的 C 推理工程也给了app_yolov5.cpp、yolo.cpp、sampleOptions.cpp 再加上 Dockerfile基本是解压、构建镜像、跑脚本、出引擎。适合正在搞树莓派、RK3568 或者 Jetson 部署的人也适合只想用成熟流程快速拿到一个能跑的 INT8 模型的人。这里先说结论这套东西值得下载但你得先弄清楚剪枝的语义和 TensorRT 的脾气否则很容易跑通却没有效果。2. 剪枝不是玄学先从通道稀疏和权重稀疏里选对路子2.1 结构化剪枝 vs 非结构化剪枝差在哪很多人一提剪枝就以为是“把不重要的权重清零”其实这是两回事。非结构化剪枝也叫权重稀疏是把权重值直接置零模型文件会变小但计算图没有变推理时仍然要访问那些零值内存除非后端有专门的稀疏矩阵算子。TensorRT 对权重稀疏的支持很有限如果你只做非结构化剪枝转 engine 之后会发现延迟一点没降白忙一场。结构剪枝则是直接把整个卷积核或者特征通道删掉特征图维度真的变小了后续所有层计算量都会下降TensorRT 对这种通道级变化非常敏感是部署端真正能吃到的剪枝收益。YOLOv5 的 C3 和 detect head 里都有大量卷积层实际工程里很少有项目只用一种。常见组合是先用结构化剪枝把通道数砍下来模型体积明显下降再用权重稀疏把精度补稳。摘要里说的“减少 70% 以上”绝大部分来自结构化剪枝权重稀疏贡献的是精度补偿。这个顺序不要反一上来就做权重稀疏后面再想转成通道剪枝基本要重新训练。选结构化剪枝时重点看 FPN 层和 head 处的 1x1 卷积通道。这些层直接影响多尺度特征融合和分类置信度剪多了小目标直接消失。所以不能拍脑袋给一个全局剪枝率需要对每一层做敏感度分析看哪层剪掉以后精度掉得少哪层碰都不能碰。压缩包里如果带了 config 文件要用的就是一张 per-layer 剪枝率表而不是一个简单的 0.7。2.2 从压缩包里的文件名看懂你的剪枝产物要过哪几道手我把这个压缩包里的文件先翻译一遍你就大概明白它是什么定位了。文件作用Dockerfile容器环境定义避免 TensorRT/CUDA 版本打架setup.cfgPython 包配置让本地脚本能用一致路径 importcommon.cmakeCMake 公共编译配置sampleOptions.cpp命令行参数解析控制推理 batch、输入路径等yolo.cpp / utils.cppYOLO 后处理和工具函数app_yolov5.cppTensorRT 推理主入口logger.cppTensorRT 日志回调kernel_function.cu自定义 CUDA 算子补偿 TensorRT 不支持的 OP看到 kernel_function.cu 就应该明白剪枝量化后的模型大概率是导出成 ONNX再转成 TensorRT engine 的。某些算子在 ONNX 里合法但 TensorRT 的 INT8 engine 不支持这时候就需要用自定义 CUDA kernel 去补。如果你只是想拿着一个 .pt 文件直接测精度这个包不是给你用的它的核心是部署链路而不是训练或者调参。setup.cfg 是我会优先看的东西。很多剪枝脚本都要求把自己当本地包安装缺了这一步跑python xxx.py的时候会报各种 import 找不到。所以下载完后第一件事不是找模型而是看 setup.cfg 里的 packages 和 name 字段确认你安装的模块名和你脚本里 import 的一致。common.cmake 则是给后续 C 工程用的剪枝量化产物最终要交给 app_yolov5 加载这个 CMake 不配好后面编译会非常痛苦。2.3 剪枝比例怎么定先跑一遍精度基线再动刀我一般会在动剪枝前先跑一次原始模型的 mAP作为基准分数。然后在 0.3、0.5、0.7 三档剪枝比例下各做一次验证画一条“剪枝比例 vs mAP”的曲线。很多刚上手的人直接拿 0.7 一剪精度掉了一截还找不到原因。实际上 detect head 的通道是最后输出的关键主流做法是主干多剪、head 少剪甚至完全不剪 head。剪枝代码里如果有配置文件你要找的是一张 per-layer 的剪枝率表而不是一个全局的浮点数。下面是一段最容易落地的评估思路。假设你已经有一个训练好的 checkpoint验证剪枝前后的差距import torch import torchmetrics def evaluate(model, val_loader): # 用 torchmetrics 计算 mAPYOLOv5 官方在验证阶段也用类似指标 metric torchmetrics.detection.mean_ap.MeanAveragePrecision() model.eval() with torch.no_grad(): for images, targets in val_loader: preds model(images) metric.update(preds, targets) return metric.compute() # 先备份原始模型并得到 baseline base_model torch.load(yolov5s.pt) base_map evaluate(base_model, val_loader) # 在多个剪枝比例下做对比这里只是示意 for ratio in [0.3, 0.5, 0.7]: pruned_model prune_model(base_model, ratio) temp_map evaluate(pruned_model, val_loader) print(fprune ratio {ratio}, mAP {base_map} - {temp_map})这段代码背后的逻辑是先拿到 baseline再去遍历剪枝比例。注意prune_model必须保证返回的模型 state_dict 的 key 名字和原来的能对上否则后面继续做量化感知训练时优化器和 BN 层的参数会错位。参数层面ratio 的语义在不同脚本里不一样有的地方 0.7 表示保留 70% 通道有的地方表示剪掉 70%。打开 config 之前先确认我就在这上面翻过车。还有一点稀疏训练时 BN 层的 gamma 值会变得很小这是剪枝时的重要依据。你可以把剪枝前的 BN gamma 打印出来看分布是否明显分层如果大部分值都在 0 附近说明稀疏训练没到位剪枝阈值很难选准。3. 一键运行落地Docker 环境、剪枝脚本与 TensorRT 编译3.1 先读懂压缩包里的文件组成这个压缩包并不是只有一个训练脚本而是把从环境到推理全链路都打包了。我整理了一张表你按图索骥就好文件作用使用阶段Dockerfile容器环境定义装 PyTorch、TensorRT、CUDA环境构建setup.cfgPython 包安装配置环境构建common.cmakeCMake 公共配置C 编译sampleOptions.cpp命令行参数解析推理启动yolo.cpp / utils.cppYOLO 后处理与工具函数推理app_yolov5.cpp主推理程序推理logger.cppTensorRT 日志回调构建/推理kernel_function.cu自定义 CUDA 算子构建这里最容易被忽略的是 setup.cfg。它能让剪枝脚本被当作本地包安装避免 Python import 路径各种打架。我在别的项目里经常遇到脚本能跑但自写模块导入失败就是因为没先做这一步。如果你解压后看到 setup.cfg第一件事就是在容器里执行pip install -e .把当前路径注册成可导入包。3.2 Docker 环境怎么搭不是上来就装 PyTorch压缩包里的 Dockerfile 建议优先使用不要在自己机器上手动装 TensorRT。TensorRT 版本依赖非常硬CUDA、cuDNN、TensorRT 三者必须对齐稍微错一个版本builder 就报各种找不到符号。用 Docker 是避免环境灾难的最省事方式。解压后先构建镜像# 解压压缩包 tar -zxvf yolov5-prune-quant.tar.gz cd yolov5-prune-quant # 构建镜像 docker build -t yolov5-prune-quant . # 进入容器并映射当前目录 docker run -it --gpus all --ipchost \ -v $(pwd):/workspace \ yolov5-prune-quant /bin/bash这段命令的关键在这里--ipchost是为了让 PyTorch 的 DataLoader 多进程时能使用共享内存不加这个训练或验证可能卡在 dataloader 的wait()上。-v $(pwd):/workspace是把当前目录映射到容器里保证剪枝脚本生成的权重文件直接落在宿主机免得到时候还得从容器里docker cp出来。--gpus all需要 NVIDIA Container Toolkit 支持先确认docker run --gpus all能正常运行再继续。进入容器后第一件事不是直接跑脚本而是验证 PyTorch 是否真的能用 GPUpython -c import torch; print(torch.cuda.is_available())如果输出True说明环境正常。这一步不是废话很多镜像里的 PyTorch 编译时没开 CUDA所有操作都在 CPU 上跑剪枝 70% 的模型可能在一次验证集遍历中就让你等到怀疑人生。3.3 剪枝量化的一键脚本参数怎么改进入容器后压缩包里的剪枝脚本一般会接受一个配置文件或者直接在命令行传参。我拿一段最常用的做法示意具体文件名以你解压后的实际内容为准# 第一步稀疏训练给 BN 层加 L1 约束 python yolov5_prune.py --data coco.yaml --weights yolov5s.pt --sr 0.001 # 第二步用稀疏化权重做结构化剪枝 python yolov5_prune.py --data coco.yaml --weights runs/train/exp/weights/last.pt --prune 0.7 # 第三步导出 ONNX python export.py --weights runs/train/exp/weights/pruned.pt --include onnx --dynamic这里--sr是稀疏率控制 BN 层 gamma 的 L1 正则强度。常用值在 0.0005 到 0.002 之间。太大会导致 BN 缩放因子一起被压掉通道全部消失精度暴跌太小又达不到稀疏效果剪枝时阈值选不准剪出来的模型参差不齐。--prune 0.7通常表示剪掉 70% 通道保留 30%这个语义要和脚本里的实现确认。--dynamic是保留动态 batch 维度这样后续 TensorRT 构建 engine 时能适配不同的 batch 大小。这一步最常见的翻车是稀疏训练完成后没有做 BN 重归一化剪枝脚本按全局阈值切通道时有的层保留了一些本应去掉的通道有的层被剪秃。所以脚本里通常会打印每层保留的通道数。你看到任何关键层通道数变成 1就可以停下来了那基本是废模型不要再往下量化。3.4 编译 C 推理程序app_yolov5 使用示例剪枝量化后的 ONNX 要转成 TensorRT engine然后用压缩包里的 C 工程去加载推理。编译方式依赖 common.cmake常规流程如下cd /workspace mkdir build cd build cmake .. -DTensorRT_ROOT/usr/local/tensorrt make -j8 # 产出可执行文件 app_yolov5 ./app_yolov5 --enginepruned_int8.engine --inputsample.jpgcmake 关键参数是TensorRT_ROOT容器里安装的 TensorRT 不一定在默认路径需要确认。如果kernel_function.cu编译报错多数是 CUDA 架构编号不对常见解决办法是在 CMake flags 里加-gencodearchcompute_75,codesm_75或者sm_80对应不同显卡架构。编译通过后先用一张小图跑通再上批量验证。sampleOptions.cpp里一般会处理--calib、--batch、--workspace等参数具体字段名以源码为准。总之这个 C 工程是最后落地的门面剪枝量化最终效果好不好得看它加载 engine 后的推理结果。4. 量化感知训练和 TensorRT INT8精度保持的关键细节4.1 量化感知训练别等转完再后悔很多项目拿浮点模型直接转 INT8结果 mAP 掉了四五个点才到处找原因。量化感知训练的核心是前向模拟量化把权重和激活转成低比特再转回浮点让网络在训练阶段就看到量化噪声。剪枝之后模型结构已经被改动直接做训练后量化风险非常大所以这个压缩包里的一键流程会包含 QAT。你在跑的时候要留意训练日志里是否出现了fake_quant或者observer相关输出有就说明量化分支已经生效。QAT 里的关键参数是量化位数和 per-channel / per-tensor 的粒度。大多数 TensorRT 部署场景卷积权重用 per-channel 效果更好因为不同通道的权重值域差异很大激活值一般用 per-tensor动态 per-channel 的开销太大。实际使用中我建议先用 per-tensor 跑一版看精度差多少如果差超过 1 个点就改成 per-channel 试试。很多新手一上来就追求 per-channel但 QAT 训练时间会明显变长收敛也慢得不偿失。4.2 TensorRT 的 INT8 引擎参数设置当你要把 QAT 后的模型转成 TensorRT engine 时app_yolov5.cpp 或 export 脚本里会看到类似的 builder 配置// 创建 builder 和 network auto builder createInferBuilder(sampleLogger); builder-setInt8Mode(true); builder-setInt8Calibrator(calibrator); builder-setMaxBatchSize(32); builder-setWorkspaceSize(1 30); // 1GB 工作空间 // 对每个需要的层设置精度 auto layer network-getLayer(i); layer-setPrecision(DataType::kINT8); layer-setOutputType(0, DataType::kINT8);这段代码里最容易踩坑的是setInt8Calibrator。如果你用的是训练后量化TensorRT 需要校准数据集生成激活分布如果你用的是 QAT网络里已经带了 fake_quant 信息TensorRT 会尝试直接沿用不用再走完整校准流程。两种情况的 calibrator 不能混着用。setWorkspaceSize也不是越大越好虽然大 workspace 能让 TensorRT 尝试更多层融合但显存不够时 builder 直接崩。移动端部署建议 1GB 以下桌面 GPU 可以给到 2GB。setMaxBatchSize这个值不是运行时 batch 就必须相等。TensorRT dynamic shape 模式下batch 的范围由优化 profile 决定。如果你构建时固定了 maxBatch运行时只能小于等于它否则需要重新构建 engine。工程上通常设一个上限比如 32避免每次改 batch 都重建。4.3 校准数据集精度和速度的平衡点如果走训练后量化校准数据一般取 500 到 1000 张有代表性的图片。太少则激活分布不够太多则校准时间太长。选择校准图时注意不要只用白天场景否则夜间目标全被截断。你可以把验证集按目标大小做分层采样确保大中小目标都有。还有一个容易被忽略的点校准数据必须经过同一套预处理包括 letterbox、归一化、颜色通道顺序。很多项目精度翻车就是因为输入图片 BGR/RGB 顺序没统一。量化后的评估通常关注 mAP0.5 和 mAP0.5:0.95 两个指标。建议以 mAP0.5 作为稳定性指标它更接近用户实际看到的效果。如果 INT8 比 FP16 掉点超过 1%先检查校准数据和预处理再考虑是否需要混合量化。上一章提到的 kernel_function.cu 在这里就派上用场了。当某些层在 INT8 下精度损失过大你可以单独把那几层指回 FP16TensorRT 的setPrecision是按层控制的不是只能全局一刀切。5. 避坑指南剪枝量化到 TensorRT 部署的五个翻车现场5.1 剪枝后模型尺寸变小但推理延迟没降现象模型文件小了 70%TensorRT 跑起来帧率没变甚至更慢。原因非结构化剪枝产生的稀疏矩阵在 TensorRT 里没有对应 kernel零值多但计算依旧稠密或者结构化剪枝只剪了 BN 层没有真正删掉对应卷积通道。解决在剪枝脚本里确认输出模型是做过 compact 处理的也就是把被剪通道从卷积权重中真正移除导出 ONNX 后可以画一下网络结构看通道数是否真的减少了。如果只是 mask 层那就要回到结构化剪枝重新处理。5.2 INT8 engine 精度掉到没法看现象mAP0.5 从 0.8 掉到 0.5。原因没有做量化感知训练直接拿剪枝模型做训练后量化或者校准数据集数量太少只有几十张。解决确认一键脚本里是否包含 QAT 阶段把校准图数量加到 500 张以上且采样策略覆盖各类场景对敏感层单独开 FP16 混合精度。还有一个细节是校准前要把模型的 BN 层冻结尽量避免在校准过程中参数被改变。5.3 TensorRT builder 报 layer 不匹配现象构建 engine 时提示 TensorRT 无法匹配 layer 输入或者维度对不上。原因剪枝后 ONNX 里某些层输出通道为 0或者动态轴没传对。解决用onnx.checker检查模型确保所有卷积层输出通道大于 0在剪枝脚本中过滤掉保留 channel 为 0 的层。也可以把 ONNX 加载进netron里查看可疑层的输出维度定位更直观。5.4 CMake 编译 kernel_function.cu 报错现象编译时找不到 CUDA 头文件或者 nvcc 报架构不支持。原因容器里 CUDA 路径没加或者 GPU 架构过新/过旧和编译参数不匹配。解决在 CMake 命令行指定-DCUDA_TOOLKIT_ROOT_DIR/usr/local/cuda在 nvcc flags 里加正确架构编号比如-gencodearchcompute_75,codesm_75。如果是 RTX 30 系以后的卡可能要sm_86或者sm_89不要照抄别人的注释。5.5 Docker 里能跑出容器就找不到库现象构建好的 app_yolov5 放到其他机器运行提示libnvinfer.so.7找不到。原因TensorRT 是动态库链接的容器外没有同样版本路径。解决用ldd app_yolov5查看依赖把相关 so 复制到可执行文件目录或者用LD_LIBRARY_PATH指定路径。更省心的做法是构建一个 runtime 镜像只装 TensorRT 运行依赖把可执行文件放进去这样就不会受宿主机环境干扰。6. 进阶验证技巧用 app_yolov5 的 CLI 参数做端到端回归剪枝量化做完最担心的就是精度掉点。我自己的习惯是拿同一组测试图片分别跑原始 FP16 引擎和剪枝量化后的 INT8 引擎然后对比输出结果。压缩包里的 sampleOptions.cpp 已经帮你把命令行参数留好了核心用法是# 用同一批测试图分别跑两个引擎 ./app_yolov5 --engineyolov5s_fp16.engine --input./val2017 --output./pred_fp16 --batch16 ./app_yolov5 --enginepruned_int8.engine --input./val2017 --output./pred_int8 --batch16跑完之后用一段简单的 Python 脚本读取两个输出目录的 json 或 txt 结果算一下 mAP 的差值和单个目标的置信度分布。如果 INT8 和 FP16 的 mAP0.5 相差在 1% 以内这个剪枝模型就可以上线。注意 batch 参数要设为一样否则前后处理逻辑会影响结果对比。我还会额外记录三组数据模型体积、单帧耗时、mAP 差。把它们做成一张小表引擎体积单帧耗时mAP0.5原始 FP1647MB8.2ms0.812剪枝 INT813MB4.1ms0.803这样我能一眼判断这次压缩是不是值得。如果剪掉 70% 体积精度只掉不到 1%那当然划算如果体积只小了 20%精度掉 3%那说明剪枝策略和量化配置有问题需要回到前面重新调参数。从那以后我每次做剪枝量化都强制自己先把原始 FP16 引擎导出再跑一遍这个对照流程不拍脑袋直接上线。这个包真正的价值不是点一下就出结果而是让你在一套固定流程里快速试错找到自己业务场景下的压缩边界。希望帮到你。本文还有配套的精品资源点击获取
