多彩编程 多彩编程MZPH · CODE BLOG
ARTICLE DETAIL

文章详情

深耕前端与后端开发技术的一线实战笔记与踩坑复盘。

YOLOv5转ONNX实战:导出链路、避坑排查与部署衔接

YOLOv5转ONNX实战:导出链路、避坑排查与部署衔接 简介面向YOLOv5模型跨框架部署需求这份资源聚焦PyTorch到ONNX的转换全流程适合目标检测开发者、算法工程师以及需要将模型部署到移动端或边缘设备的工程人员。内容结合YOLOv5自身结构梳理了模型导出前的权重准备、动态输入张量设定、导出后的结构校验与精度核对同时给出借助ONNX Runtime进行图优化和加速推理的实用参考还提醒了不同框架间算子映射差异可能带来的精度损失与调整思路。资源包为zip格式整体大小约1.02MB当前展示的文件总数与类型明细暂未同步实际内容以解压后为准。已有1342人浏览学习适用于希望打通PyTorch训练与ONNX推理链路、并进一步转换CoreML或TFLite格式的读者。通过该资源可快速建立清晰可靠的转换路径减少踩坑成本为后续跨平台上线打下基础。1. 从 PyTorch 到 ONNXYOLOv5 部署前绕不开的转换关把 YOLOv5 模型转换到 ONNX已经是前端训练和后端推理之间的固定工序你用自有数据集把 best.pt 跑出来到了树莓派、RK3568 或 Jetson 这类设备上PyTorch 运行时要么装不进去要么慢得没法用换成 TensorRT、RKNN、OpenVINO 这些引擎之前普遍的做法是先导出 ONNX 这个中间格式。卡在这道工序上的人不少常见现象是命令执行完但 onnxruntime 加载报错或者框能画出来坐标却和 PyTorch 差一截。下面我会把从 PyTorch 到 ONNX 的转换经验讲清楚从 export 链路、版本差异、参数语义到排查思路照着做你就能拿到能用的 .onnx。2. 先搞清楚导出链路export.py 做了什么以及版本差异2.1 导出链路拆解从权重文件到 ONNX 图的四步先别急着敲命令。很多翻车其实在动手之前就注定了因为不理解 export 脚本究竟把模型怎么了。整个导出流程可以拆成四步来看。第一步是加载权重。export.py 内部会按 torch.load 的方式把 .pt 里的模型结构、state_dict 和训练元数据一起取出来然后把 EMA 权重合回主模型。这一步常见的坑是权重文件的 PyTorch 版本和当前环境对不上比如用 PyTorch 2.0 保存的模型在 1.10 环境里 load 直接报错这类问题不属于 ONNX但会最先拦住你。第二步是切推理模式。模型要执行 model.eval()同时把 BN 层冻结。YOLOv5 在训练时 BN 的统计量会随 batch 变化如果不换到 eval导出的图里 BN 语义是错乱的跑出来的置信度全部失真。这一步 export.py 已经做了但你要是自己写导出脚本很容易漏掉漏掉之后的现象非常隐蔽模型能跑分数全飘还不报错。第三步是追踪导出。YOLOv5 走的是 torch.onnx.export 的 tracing 路线也就是给模型喂一个固定 shape 的假输入把 forward 实际执行到的算子逐个记录下来映射成 ONNX 节点。trace 模式有两个直接后果输入尺寸在导出那一刻就被钉死python 里的 if、for 这类控制流会被拍平。这就是为什么 export 必须显式给 --img 和 --batch也是为什么动态 shape 不能靠改代码实现必须依赖 --dynamic 参数。第四步是校验与简化。导出结束会调用 onnx.checker 做结构检查然后在有 onnxruntime 时试着跑一次推理。之后如果加了 --simplify会调 onnx-simplifier 做常量折叠和冗余节点清理把图变小。这一条链路记住后面所有排查都能对上号。2.2 版本差异Focus、SPPF 与 SiLU 在导出时的表现很多人在导出时遇到算子不支持其实和 YOLOv5 版本强相关。我经手的项目里有 v5.0 的老仓库也有 v6.0、v7.0 的新仓库导出的图长得完全不一样。v5.0 的 backbone 第一层是 Focus它的做法是把输入按通道切片重排等效于一个 6x6 卷积但没有标准卷积算子结构。ONNX 图里会体现成大量 Slice、Concat 节点。问题在于Focus 的切片索引在 ONNX 标准里没问题但落到部分 NPU 编译器和老版 TensorRT 上这些 Slice、Concat 的处理效率极低甚至直接拒绝编译。v6.0 之后 Focus 被重构成一个 stride2 的 6x6 卷积导出图干净很多这也解释了为什么同等精度下 v6.0 的 onnx 更好部署。SPP 到 SPPF 的变化同样影响导出。v5.0 用的是 SPP多个不同 kernel 的 MaxPool 并行v5.0 中期之后换成 SPPF一个 kernel 的 MaxPool 串行三次再拼接。串行池化在 ONNX 里就是几个 MaxPool 节点连在一起计算量更小但有些 NPU 工具链对连续 MaxPool 的折叠处理不好转 RKNN 时偶尔会报不支持。遇到这种情况通常解法不是换结构而是在导出后手动把几个池化合并或者接受该部分使用 CPU 算子执行。激活函数是另一个点。v6.0 起 YOLOv5 默认用 SiLU也就是 swish。在 PyTorch 1.9 及更早的版本里SiLU 到 ONNX 的映射不完整opset 11 下经常直接报 Exporting operator silu not supported。新版本 PyTorch 配合 opset 12 以上就没有这个毛病了。建议直接把环境升到 PyTorch 1.11 以上再看。这一节讲的三个差异其实就是你看网上搜到的 yolov5 网络结构图时那些结构图变化背后的原因。把差异整理成一张表决定导出前先确认自己属于哪一行结构/激活v5.0v6.0 及以后导出与部署影响Focus切片通道拼接stride2 的 6x6 Conv老版导出的 Slice/Concat 在 NPU/TensorRT 上易踩坑SPP/SPPFSPP 并行池化SPPF 串行池化连续 MaxPool 在部分编译器需要折叠激活LeakyReLUSiLUSiLU 需要 opset12旧 PyTorch 会直接报错2.3 环境匹配PyTorch、onnx 与算子集的兼容关系这部分用一条命令就能验证在终端敲python -c import torch, onnx, onnxruntime; print(torch.__version__, onnx.__version__, onnxruntime.__version__)。我一般要求环境满足 PyTorch 1.10 以上、onnx 1.12 以上、onnxruntime 1.12 以上。注意导出过程不依赖 CUDACPU 环境一样可以导出常见的 cuda 和 pytorch 匹配问题只影响你加载权重做验证的速度不影响 ONNX 正确性。这点很多人误解在 GPU 服务器上导出失败就以为是显卡驱动坏了其实是算子集或者 PyTorch 版本问题。另外PyTorch 2.x 用户会遇到一个新情况torch.onnx.export 默认走了 dynamo 路径某些老模型的 trace 会卡住或者导出的图结构和你预期不符。我的处理办法是在调用 export 时显式设置 dynamoFalse回到旧版稳定路径图的结构和 1.x 时代完全一致。这在社区里已经是很常见的操作了不需要改任何模型代码。再看算子集。onnx 的 opset 是协议版本号数值越高支持的算子越多但下游引擎支持度往往滞后。导出只要高于某个下限就行不是越高越好。常规推荐是 opset 12兼顾 SiLU 支持和老设备兼容。你要是目标平台是 RKNN先试 11目标是 TensorRT 8.x12 或 13 都稳。这个选择我下一章展开讲。理解链路、版本、环境三件事后再去看导出命令就不会是一堆黑匣子参数了。3. 跑通最小转换一条命令导出 ONNX再拆解每个参数3.1 最小导出命令与参数拆解先给你一个能跑的最小命令用官方 export.py 在训练完之后的工作目录执行python export.py \ --weights runs/train/exp/weights/best.pt \ --include onnx \ --img 640 \ --batch 1这条命令会在 best.pt 旁边生成 best.onnx。这里拆开说每个参数。--weights 指向训练产生的权重我习惯用 best.pt 而不是 last.pt因为 best 的 mAP 通常高几个点部署后线上效果更直观--include 控制导出产物torchscript、onnx、engine 都可以用逗号拼接只在确定要 TensorRT 时才加 engine因为每次跑会调 trtexec很慢--img 是输入边长YOLOv5 支持一个值表示正方形也支持两个值分别表示高和宽实际生产里 640 是性价比很高的选择低于 320 小目标基本丢光高于 1280 推理时间翻倍--batch 是固定批量导出时显式写成 1 最稳因为 batch 维度进图之后一旦变动后面做量化或转 TensorRT 都要重新处理。还需要注意一个细节export.py 末尾会顺带打印模型的参数量和 GFLOPs并且如果装了 onnxruntime会自动调用它做一次结构校验。看到 No errors detected 这一行说明图至少是结构完整的但结构完整不代表精度一致后者我放在第四章讲。3.2 三种导出形态裸输出、解码输出与端到端这是转换里决策价值最大的一步因为同一个模型可以导出成三种不同形态下游处理逻辑完全不同。默认不加任何标志导出的是裸输出相当于把 Detect 层之前的 85 维特征图直接给你输出是三张类似 [1, 3, 80, 80, 85] 的 tensor85 对应 cx、cy、w、h、objectness 和 80 个类别得分。所有解码都要自己做这是自由度最高、但最容易出错的形式。第二种是用 --grid 把解码逻辑编进图。加了之后导出的图上带 anchor 网格解码、sigmoid 和坐标换算输出直接是 [1, 6300, 6] 形状的检测框列表6 列是 x1、y1、x2、y2、score、class。这个形态最推荐省掉自己写 decode而且 anchor 会以常量形式固化在图里不会发生后处理用了错误 anchor 导致的错位问题。代价是输出 6300 个候选框还得在外层做 NMS。第三种是 --end2end把 NMS 一并编进图。YOLOv5 的官方实现基于 torchvision.ops.nms导出的 ONNX 会包含 NonMaxSuppression 算子输出变成定长的检测结果。这个形态看起来最美但落地时约束也最多NMS 这个算子在老版本 onnxruntime 和很多 NPU 工具链里不支持转 TensorRT 也要 8.2 以上才稳。我的建议是优先第二种等下游引擎验证了 NMS 算子能力再考虑第三种。end2end 导出的参数还有 --conf-thres、--iou-thres 和 --max-det它们只在 end2end 模式下生效普通导出写了也不报错但并不会约束输出。如果选择裸输出我给出一个可以抄走的解码核心段作用是验证导出的图至少算得对。以单张 640x640 输入为例假设已经用 onnxruntime 拿到三个特征图import numpy as np def decode(feats, anchors, stride): # feats: (B, 3, Ny, Nx, 85)anchors 与 stride 按缩放级别一一对应 B, A, Ny, Nx, C feats.shape # 构建网格偏移索引方式要和模型对齐iy 对应高ix 对应宽 grid np.stack(np.meshgrid(np.arange(Ny), np.arange(Nx), indexingij), axis-1) # 中心坐标sigmoid 后的值乘 2 减 0.5再加网格偏移最后乘 stride feats[..., 0:2] (feats[..., 0:2] * 2 - 0.5 grid) * stride # 宽高sigmoid 后的值乘 2 再平方最后乘对应 anchor feats[..., 2:4] (feats[..., 2:4] * 2) ** 2 * anchors return feats.reshape(B, -1, C)这段代码的核心就两行换算中心点坐标要加 grid 偏移再乘 stride宽高要乘上对应尺度 anchor并且 v6.0 之后的公式里有减 0.5 这个偏移量老版本代码没有。写这个的目的不是让你在正式项目里手工解码而是遇到输出对不齐时能自己动手定位不至于把整个流程当黑匣子。3.3 opset 与 simplify11、12、17 怎么选上一节说 decode 是决策点这一节说第二个决策点算子集。先给表格opset支持的典型算子适合场景需要小心的点11基础 Conv、MaxPool、SliceRKNN、老 TensorRT 7SiLU 映射不完整12原生 SiLU、NMS 支持齐全多数部署场景的默认值老 NPU 工具链偶有不识别13-17新注意力算子、动态 shape 支持更好最新 onnxruntime、TensorRT转 RKNN 或 TVM 时失败率高我默认选 12理由很简单SiLU 原生可用end2end 的 NMS 也在这个级别被广泛支持而且 onnxruntime 和 TensorRT 8.x 都消化得很好。只有在明确目标平台是 RKNN 这一类资源受限工具链时才会主动降到 11因为它们的解析器更新慢高版本算子集容易直接报 unsupported。再讲 --simplify。onnx-simplifier 做的事是常量折叠、冗余形状节点清理以及把一些可以合并的算子合并起来。导出后的图经过它体积通常能小三分之一在树莓派这类设备上加载更快。但简化不是零风险它偶尔会把动态 shape 相关的计算节点错误折叠导致运行时 shape 推导出错。我自己的习惯是保留一份没简化过的 onnx 做备份简化后的版本先跑一次 onnxruntime 校验这个后悔药很重要尤其是还要把 onnx 喂给后续量化工具时未简化版本反而更容易排查。所以你第一次上手做转换我的建议是先用裸命令跑通最小导出再加 --grid需要时再上 --simplifyopset 定 12。别一上来就 end2end那只是看起来省事。4. 转换避坑从报错到输出对不齐的排查经验4.1 导出直接报错算子不支持与版本错配现象是命令跑到一半终端打出类似 Exporting operator silu to ONNX opset version 11 is not supported 就停了。第一次遇到的人会以为是模型坏了其实不是。原因是 PyTorch 1.9 及之前版本对 SiLU 的 ONNX 映射不完整或者你用了 opset 11 而该算子还没进标准集。解决办法按优先级来把 --opset 改成 12 或 13 重导如果还不行把 PyTorch 升到 1.11 以上再不行检查是不是自己改了模型结构把 SiLU 换成了别的自定义激活。每次改完重新导出不要在原图上反复折腾。另一个常见报错是 onnx.checker 过了但 onnxruntime 加载时报 Unsupported ONNX ops。现象是报出某个冷门算子名。这种情况常见于用了官方仓库之外魔改过的检测头比如加了 DeformableConv 或者自定义注意力。原因是这些自定义算子没有对应的 ONNX 实现。常规做法是看这个算子能不能透传不能的话就得把这层改动移到图外即在导出前把模型里这层替换成等价的标准算子组合。这个属于吃力但必要的活没有捷径魔改越深导出越疼。还有一种间接报错torch.load 阶段就挂了提示版本不匹配。现象是报错里带 weights generated by incompatible version of PyTorch。原因是 .pt 文件本身就是版本绑定和 ONNX 一点关系都没有。解决思路很简单要把权重落到一个中间格式比如先在同版本 PyTorch 里转存或者直接在当前环境重跑一次训练。这种事我碰过两次之后现在所有导出环境都固定 PyTorch 版本不跟着最新走。4.2 检测框对不齐九成是解码错了不是转换错了现象很典型同一张测试图torch 推理和 onnx 推理都能画框但 onnx 这边的框整体偏移小目标尤其飘或者置信度全部偏低。网上搜yolov5 转 onnx 后检测不准能搜出一堆但大量是解码错配导致的翻车现场。原因是导出链路没问题问题在推理端的解码没有对模型结构。YOLOv5 有两套网格解码逻辑一套是 v6.0 之后基于 anchor 的新实现坐标公式里带减 0.5 的偏移一套是老版本实现归一化缩放不一样。如果你导出的是裸输出又用网上抄来的老版 decode 代码那 v6.0 模型出来的坐标就会差出几个格子。解决建议也是最省事的导出时带 --grid。这样 anchor 和网格偏移直接以常量编进图里onnx 输出就已经是解好的框不存在解码错配。如果你坚持裸输出就不要用第三方轮子直接抄你导出那个版本仓库里 Detect 前向的公式一行一行对着改。还有一种判断方法导出前后分别用官方 val.py 跑同一份验证集。如果 mAP 差距超过 0.5%大概率是解码或后处理问题回头查代码1% 以内属于可接受的工程误差。注意mAP 是 NMS 之后的指标所以差异里既有解码贡献也有 NMS 参数贡献。别一上来就怀疑量化先确认这条链路。4.3 动态尺寸与多 batch能导出来不一定跑得起来我先说结论YOLOv5 的动态导出我的建议是只开 batch 维度的动态长宽维度永远固定。现象是有人用 --dynamic 成功导出但上线时换了个 320x320 的输入onnxruntime 直接报 Input image of size is invalid。原因有两层第一YOLOv5 的下采样倍率是 32模型内部每层输出尺寸必须是 32 的倍数320、352、416、480、544、608、640 这些值没问题321 或 340 都会崩第二SPPF 里串行池化在非整除尺寸上算子图处理起来会有边界不齐。多 batch 也有类似情况。--batch 4 导出的模型在推理引擎里一旦遇到 batch1 的请求一些后端优化会重置 buffer导致性能骤降或者直接报错。所以我一般在生产环境固定 batch1用并发去扛吞吐而不是图里留动态 batch。真需要多尺寸就在 32 的倍数里做几个固定档位导出多份 onnx或者用 --dynamic 但把 shape 范围测试清楚再上。YOLOv5 的导出脚本虽然给你动态选项但实际落地时固定尺寸永远是省心方案。4.4 int8 量化与 RKNN 转换的精度崩坏现象是最扎心的fp16 的 onnx 部署得好好的一心想把 onnx 量化成 int8 跑快点结果小目标全丢mAP 掉二十个点转 RKNN 之后更离谱模型直接输出一堆空列表。这个我在 RK3568 上反复踩过。原因是多方面的但第一个元凶几乎都是量化方式不对ONNX Runtime 自带的 dynamic int8 量化按激活的绝对范围算 scale而检测头的输出经过 sigmoid 后集中在 0 到 1 之间绝对范围很小量化步长把细节全抹掉了。正确的量化流程是离线校准准备 200 到 500 张代表图片在导出时把检测头部分保持 fp16backbone 和 neck 做 per-channel int8然后由校准工具统计每一层的激活分布来定 scale。具体到 RKNN 就是 rknn-toolkit2 里用 dataset.txt 喂校准图片config 里把 quantized_dtype 设为 asymmetric_quantized-8检测头如果量化后精度崩就把 detect 层拆出去不量化。还有一个容易被忽略的点量化前不要用 --simplify 之后的图。simplify 会合并掉一些 BN 折叠需要的分支导致校准时的激活分布和真实模型对不上。备份原始 onnx 这个习惯在这一步救过我两次。5. 导出了怎么验证onnxruntime 对比与下游衔接5.1 用 onnxruntime 复现前向并对比输出导出完成别急着上设备先在本地用 onnxruntime 跑一遍。最省事的验证是直接用官方 val.py 对 onnx 跑验证集python val.py --weights best.onnx --data your_dataset.yaml。它内部会自动调用 onnxruntime并把结果和训练时的 mAP 对比你只需要记住训练时的数字。想更细一点可以手写一段对照import numpy as np import onnxruntime as ort x np.random.rand(1, 3, 640, 640).astype(np.float32) # 建议换成真实图片 eng ort.InferenceSession(best.onnx, providers[CPUExecutionProvider]) onnx_out eng.run(None, {eng.get_inputs()[0].name: x}) print(eng.get_outputs()[0].name, onnx_out[0].shape)这段脚本的价值在核对输出名和形状。如果你之前用 --grid输出应该是 [1, 6300, 6]如果你拿到 [1, 3, 80, 80, 85]说明导出参数没生效回去查 include。这是我最常用来抓参数错误的办法。5.2 从 ONNX 到 TensorRT / RKNN / OpenVINO 的衔接ONNX 只是中间形态真正跑起来还要看下游。TensorRT 用trtexec --onnxbest.onnx --fp16 --saveEnginebest.trtopset 12 或 13 都稳RKNN 用 rknn-toolkit2 的 rknn.config 配 mean/stdopset 11 或 12 最保险OpenVINO 用 openvino.convert_model 一行转换兼容性最好。我现在的习惯是把导出链路的 mAP 记录留档每次改动模型或导出参数后跑一次 diff这个习惯帮我抓了量化翻车也省去了后端工程师反复来问你这 onnx 到底靠不靠谱的沟通成本。希望帮到你。本文还有配套的精品资源点击获取
返回列表