YOLOv5转ONNX:模型部署关键步骤与常见坑
简介YOLOv5模型转换从PyTorch到ONNX的实操型资源包面向智能监控、自动驾驶、边缘计算等场景的计算机视觉开发者和算法工程师。资源系统梳理了将训练好的YOLOv5模型导出为ONNX的完整链路涵盖模型准备、torch.onnx.export导出参数设置、转换后正确性验证、算子兼容性排查、模型优化以及进一步迁移至CoreML/TFLite等平台的扩展操作并针对不同框架间的算子映射差异和精度损失问题给出排查思路可帮助读者理解动态图模型转为静态图格式时的关键细节从而在实际部署中少走弯路。资源采用zip压缩包封装整体约1.02MB源站暂未提供文件总数与类型明细解压后可获得与转换流程紧密相关的说明资料。目前已有1342人浏览学习适合需要快速完成模型跨平台部署或系统学习ONNX转换要点的开发者借助这份材料可减少试错成本同时为后续借助ONNX Runtime加速推理、部署到移动端或嵌入式设备打下基础。1. 把 YOLOv5 从 PyTorch 搬到 ONNX先弄清楚你为哪一步做转换YOLOv5 在 PyTorch 里训练、验证、调参都挺顺手但真要把它塞进手机 App、树莓派、RK3568 这类边缘设备或者交给 ONNX Runtime 做服务端推理PyTorch 的 .pt 权重就有点“水土不服”了。这时候把模型转成 ONNX 几乎是必经之路。ONNX 相当于一个中间格式它不绑死某个框架导出之后可以再接 CoreML、TFLite 甚至 TensorRT部署路径一下子宽了很多。这篇文章不跟你讲玄学直接按我实际拆过的流程把torch.onnx.export的参数、输出形状的变化、动态轴怎么开、NMS 后处理怎么接这些关键点过一遍。适合正在做模型部署、或者刚把 YOLOv5 训练完准备往硬件上迁移的工程师——新手能跟着步骤走熟手可以直接跳到第 5 章看坑。2. ONNX 到底改了什么动态图变成静态计算图输出形状跟着变2.1 PyTorch 模型和 ONNX 模型的本质差异PyTorch 是动态图框架模型的前向传播是“边执行边建图”你可以在 forward 里写 if 分支、for 循环甚至根据输入 shape 动态改变计算路径。这个灵活性在训练时非常舒服但在部署时就是麻烦推理框架没法预知计算图结构也就很难做内存分配和算子融合优化。ONNX 的做法是把模型“冻结”成一张静态计算图。torch.onnx.export会用一个 dummy input 实际跑一遍前向传播把沿途执行过的算子记录下来导出成一个 .onnx 文件。这意味着你在模型里写的 Python 控制流在导出时只会保留“实际走的那条分支”不会被转换成动态逻辑。这就是为什么很多人导出后发现自己模型里的某些条件判断“消失”了——不是 bug是 ONNX 的静态图特性决定的。还有一个容易被忽略的点导出的 ONNX 模型默认是“全精度 FP32”的。YOLOv5 的权重本来就是 FP32导出时一般不会主动降精度但后续你用 ONNX Runtime 推理时可以打开算子的优化选项或者用 INT8 量化把模型压到更小。这个后面专门说。2.2 YOLOv5 的 Detect 层在导出时为什么要特殊处理YOLOv5 的模型结构里Detect 层是最后一环负责把特征图解码成预测框。这个层里有 anchor 网格生成、box 坐标解码、objectness 和 class 概率计算还有一堆 reshape、transpose、sigmoid 操作。在 PyTorch 里跑没问题但导出 ONNX 时如果直接整个模型端到端导出Detect 层的解码逻辑会变成一长串算子既冗长又低效而且后续你想接 NMS 还得从这堆输出里手动抠数据。Ultralytics 官方仓库的做法是导出时用--include onnx参数但导出脚本里其实会把 Detect 层替换成一个简化版只输出原始预测张量也就是三个尺度的特征图输出不包含解码逻辑。这样导出出来的 ONNX 模型输出是三个张量每个张量的 shape 是[batch, 3, grid_h, grid_w, (5 num_classes)]其中 3 代表三个 anchor5 是 box 中心点坐标、宽高、objectnessnum_classes 是你的类别数。这里有个实际影响你在 PyTorch 里调用模型得到的输出是解码后的 boxes、scores、classes而 ONNX 模型拿到手的是原始特征图输出必须自己在外面写后处理把 grid 网格和 anchor 偏移解出来。很多刚上手的人在这一步懵了以为模型导出错了其实只是输出形式变了。2.3 export 的核心参数opset_version、dynamic_axes、input_names直接看一段我常用的导出脚本基于 YOLOv5 官方仓库改的import torch from models.experimental import attempt_load # 加载训练好的权重注意 map_location 要指定 cpu避免 GPU 显存占用 model attempt_load(best.pt, map_locationcpu) model.eval() # 构造一个 dummy input尺寸要和训练时一致 dummy_input torch.randn(1, 3, 640, 640) # 导出 ONNX torch.onnx.export( model, dummy_input, yolov5s.onnx, opset_version12, input_names[images], output_names[output0, output1, output2], dynamic_axes{ images: {0: batch}, output0: {0: batch}, output1: {0: batch}, output2: {0: batch}, } ) print(导出完成)这段代码里最关键的是opset_version和dynamic_axes。opset_version决定了 ONNX 算子集的版本版本太低可能不支持某些 PyTorch 算子版本太高某些推理引擎比如老版本 TensorRT又不认。我一般先用 12如果导出时报算子不支持再往上调。dynamic_axes声明了哪些维度是可变的这里只把 batch 维设成动态让同一个模型可以处理 batch1 也能处理 batch4。输入图像的宽高不建议设成动态YOLOv5 在固定尺寸下推理效率更高后面接 TensorRT 也更方便。导出完成后建议立刻用一个 ONNX Runtime 的 session 跑一遍确认输出 shape 和推理结果符合预期import onnxruntime as ort import numpy as np session ort.InferenceSession(yolov5s.onnx, providers[CPUExecutionProvider]) x np.random.randn(1, 3, 640, 640).astype(np.float32) outputs session.run(None, {images: x}) for i, out in enumerate(outputs): print(foutput{i1} shape: {out.shape})这一步能快速发现导出时有没有算子丢失、shape 是否对不上。注意输入张量要转成np.float32而且要按 NCHW 排布这和 PyTorch 里的 tensor 保持一致。3. 端到端导出实操准备、参数调整、后处理对接3.1 环境准备和权重文件的坑导出前先确认 PyTorch 版本和 YOLOv5 仓库版本是匹配的。YOLOv5 更新很勤不同版本的模型定义可能不一样老权重配新仓库经常出现 key 对不上的问题。我的习惯是直接用官方仓库打一个固定版本 tag比如v6.0或v7.0不要用 master 分支的最新代码否则今天能导出明天可能就报错。权重文件方面最好用训练完的best.pt而不是last.pt。last.pt是训练过程中最后一个 epoch 的权重过拟合风险高best.pt是按验证集指标挑出来的最优权重。如果你用的是官方预训练权重确保下载的是完整的 .pt 文件而不是只包含 state_dict 的权重——因为attempt_load需要完整的模型结构。有个小坑如果你训练时用了多 GPU权重的 key 里可能带module.前缀导出前需要先torch.load再手动去掉这个前缀不然模型结构对不上。常见做法是加载权重后打印一下model.state_dict().keys()看看有没有异常。3.2 导出时输入尺寸的选择640 还是训练尺寸YOLOv5 默认训练尺寸是 640x640导出时 dummy input 也建议用这个尺寸。如果你训练时用了 416 或 1280那 dummy input 要跟着改。这里有个细节导出的 ONNX 模型本身不强制输入尺寸但后续的 NMS 后处理、anchor 网格生成都是按训练尺寸设计的尺寸变了效果可能打折。所以在导出这一步就把输入尺寸定死后面的部署流程会省很多事。如果你确实需要多尺寸推理可以把dynamic_axes里的宽高维度也放开但我不推荐。放开之后模型能接收任意尺寸输入但三个输出特征图的 grid 尺寸会跟着变后处理代码必须写成通用的而且要处理宽高不是 32 倍数的输入补齐问题复杂度直接翻倍。一般情况下固定尺寸就够了。3.3 三个输出特征图怎么接解码和 NMS导出后的 ONNX 输出是三个尺度的特征图对应 stride 分别是 8、16、32shape 是[batch, 3, grid_h, grid_w, 85]假设 COCO 80 类。这里的 85 5 805 包含中心点 x、y、宽 w、高 h 和 objectness。后处理的第一步是把 x、y、w、h 从网格坐标解码回原图坐标import numpy as np def decode_output(pred, stride, anchors, num_classes80): # pred shape: [batch, 3, grid_h, grid_w, 5num_classes] batch_size pred.shape[0] grid_h, grid_w pred.shape[2], pred.shape[3] num_anchors pred.shape[1] # 把 pred 从 NCHW 转成 NHWC方便解析 pred pred.transpose(0, 1, 3, 4, 2) # [batch, 3, grid_w, grid_h, 85] pred pred.reshape(batch_size, num_anchors, grid_h * grid_w, -1) # 生成网格坐标 grid_y, grid_x np.meshgrid(np.arange(grid_h), np.arange(grid_w), indexingij) grid np.stack([grid_x, grid_y], axis-1).reshape(1, 1, grid_h * grid_w, 2) # decode box 中心点和宽高 xy (pred[..., :2] * 2 - 0.5 grid) * stride wh (pred[..., 2:4] * 2) ** 2 * anchors boxes np.concatenate([xy - wh / 2, xy wh / 2], axis-1) # [x1, y1, x2, y2] scores pred[..., 4:5] * pred[..., 5:] # objectness * class_prob return boxes, scores这个解码逻辑和 YOLOv5 源码里的Detect层一致只是把 PyTorch 张量操作换成了 NumPy。解码后得到的boxes是左上角和右下角坐标scores是每个类别的置信度接下来再做阈值过滤和 NMS。NMS 可以直接用cv2.dnn.NMSBoxes或者如果后面要接 TensorRT建议直接用 TensorRT 自带的 EfficientNMS 插件把 NMS 也并进推理图里。我自己在边缘设备上部署时更喜欢把 NMS 留在 ONNX 外面用 C 写因为 ONNX Runtime 的 NMS 算子在不同版本上行为不一致容易翻车。4. 验证导出结果比对 PyTorch 输出和 ONNX 输出4.1 全精度比对把输出拉到 CPU 逐项对导出完成后最稳的验证方式是用同一张输入图分别跑 PyTorch 模型和 ONNX Runtime对比输出的数值差异。这里要特别注意PyTorch 模型的输出是 Detect 层解码后的结果而 ONNX 输出是原始特征图两者不能直接比对。正确的做法是让 PyTorch 模型也走“不包含解码”的路径或者直接把 PyTorch 模型的输出再手动解码一次和 ONNX 的输出对比。我一般这样操作修改 PyTorch 模型的前向传播把 Detect 层替换成 Identity这样模型的输出就是原始特征图。然后对同一张图分别推理用np.allclose看数值误差import torch import onnxruntime as ort import numpy as np # 加载 ONNX 模型 session ort.InferenceSession(yolov5s.onnx, providers[CPUExecutionProvider]) # 准备输入图 img torch.randn(1, 3, 640, 640) # PyTorch 推理 with torch.no_grad(): pt_outputs model(img) # ONNX 推理 onnx_outputs session.run(None, {images: img.numpy().astype(np.float32)}) # 逐层比对误差 for i, (pt_out, onnx_out) in enumerate(zip(pt_outputs, onnx_outputs)): diff np.max(np.abs(pt_out.numpy() - onnx_out)) print(foutput{i} max diff: {diff})误差在 1e-4 量级是正常的如果超过 1e-2 就要警惕。常见原因是 opset 版本不一致导致算子实现差异或者模型里有某些算子比如 mish 激活函数在 ONNX 里的实现和 PyTorch 不完全一致。YOLOv5 的 backbone 用了 SiLU 激活老版本 ONNX 导出时会被拆成 sigmoid multiply 的组合算子精度上会有微小损失一般不影响实际效果但你要是做 INT8 量化这些微小差异会被放大。4.2 用 Netron 检查计算图结构数值验证通过后我还会用 Netron 打开导出的 .onnx 文件检查计算图结构。重点看三件事输入节点是不是只有一个images、三个输出节点的 shape 是不是和预期一致、有没有出现奇怪的算子比如ScatterND、NonMaxSuppression这种你可能没主动引入的。我遇到过一种情况导出时不小心把后处理代码写进了模型的前向传播里结果 ONNX 图里带了 NMS 算子导致模型输出完全不对。用 Netron 一眼就能看出来。Netron 是网页版工具直接把 .onnx 文件拖进去就行不用安装。检查完之后顺手看一眼模型的参数量大小正常情况下导出的 ONNX 文件应该和 .pt 权重差不多大如果相差很大说明导出过程有算子没被正确转录。5. 避坑指南版本、算子、动态轴和后处理的四类经典翻车现场5.1 报错 “Unsupported operator: aten::mish”——版本和算子不匹配现象导出到一半报Unsupported operator: aten::mish或类似的算子不支持错误。原因YOLOv5 的 backbone 用了 SiLU也叫 Swish激活函数在 PyTorch 里是nn.SiLU()导出 ONNX 时需要映射成 ONNX 的Sigmoid加乘法的组合。如果你的opset_version太低比如 9 或 10某些算子没有对应的 ONNX 映射就会报错。另外如果你用的 YOLOv5 版本很老可能用的是nn.Hardswish同样会有兼容问题。解决先把opset_version提到 12 或 13 试一次。如果还报错检查 YOLOv5 仓库是不是和你的 PyTorch 版本不匹配。我遇到过一次 PyTorch 1.10 配 YOLOv5 v6.0 导出正常换成 PyTorch 1.13 后aten::SiLU的行为变化导致导出失败最后降回 1.10 解决。所以导出前先固定 PyTorch 版本别跟着最新版跑。5.2 输出 shape 对不上——你用了动态宽高但后处理写死了现象导出时dynamic_axes开了宽高推理时输入了一张 800x600 的图结果输出特征图的 shape 和你后处理代码里写死的 grid 尺寸不匹配直接数组越界。原因后处理的解码逻辑用grid_h、grid_w生成网格坐标如果输入尺寸变了这两个值也跟着变代码里如果写死了 80x80、40x40、20x20 这种固定网格必崩。解决把后处理里的网格尺寸改成从输出张量的 shape 动态读取。或者更简单的做法导出时只开 batch 维度动态输入尺寸固定 640所有后处理逻辑按固定尺寸写这样最稳。如果你确实需要多尺寸输入那后处理必须完全参数化不要有任何硬编码的网格尺寸。5.3 ONNX 模型比 .pt 大很多——导出时把训练代码也带进去了现象导出的 ONNX 文件有几百 MB明显比 .pt 权重文件大得多。原因模型对象里包含了训练时附加的模块比如 EMA 权重、训练辅助头或者你在模型里定义了额外的属性。attempt_load加载权重时会把整个模型对象都加载进来导出时这些附加内容可能被序列化进 ONNX 图里。解决导出前手动把训练相关的模块移除只保留推理需要的部分。YOLOv5 的官方导出脚本export.py里已经处理了这些细节建议直接基于官方脚本改不要自己写导出逻辑。5.4 后处理速度比模型推理还慢——Python 循环解特征图现象ONNX 推理只要 5ms但你的后处理写的是 Python for 循环逐网格解码跑一次要 50ms整体耗时反而更长了。原因特征图解码如果用 Python 循环遍历每个 grid cell效率极低。PyTorch 代码里用的是张量并行操作你用 NumPy 重写时如果不注意向量化性能会崩。解决后处理必须用向量化操作用np.meshgrid生成网格用切片批量解码不要写 for 循环。前面第 3.3 节给的代码就是向量化写法。如果对性能要求更高把后处理挪到 C 或 CUDA 里实现Python 只做数据传输。6. 从 ONNX 到 TFLite 和 RKNN量化、形状约束与端侧验证6.1 INT8 量化先搞清楚目标硬件的量化要求ONNX 模型拿到手之后往具体硬件上搬是另一道坎。以 RK3568 为例瑞芯微的 RKNN-Toolkit 支持把 ONNX 转成 RKNN 格式但在转之前必须做 INT8 量化不然模型跑在 NPU 上的加速效果出不来。量化的核心目的是把 FP32 的权重和激活值映射到 INT8 范围模型的体积直接压缩到四分之一推理延迟也能明显降下来。RKNN-Toolkit 的量化流程一般是先加载 ONNX 模型然后用一个校准数据集跑一遍统计每层激活值的动态范围最后生成量化后的 RKNN 模型。我实际踩过的一个坑是校准数据集必须和你训练数据的分布一致否则量化后精度掉得很厉害。我见过有人拿几张网图做校准集结果模型在真实场景里检测框全偏了。校准集一般准备 200 到 500 张有代表性的图就够了。6.2 TFLite 转换时的形状约束固定 batch 和固定尺寸更好用如果你要转 TFLite 跑 Android流程是 ONNX 先转成 TensorFlow 的 SavedModel再用tflite_converter转 TFLite。这一步最容易出的问题是 ONNX 里的动态维度在 TFLite 里不支持转换直接报错。我一般建议在导 ONNX 时就把 batch 固定成 1宽高固定成 640把 dynamic_axes 全部关掉转 TFLite 时最省心。还有一个性能相关的点TFLite 对某些算子的支持不完善比如 YOLOv5 的 SiLU 激活在 TFLite 里可能被转换成多个基础算子推理速度比在 ONNX Runtime 里慢。2024 年之后 TFLite 对 mobile 端的支持已经改善了很多但转换完建议用 Android 真机跑一次 benchmark别只看转换成功就以为万事大吉。6.3 用一个小脚本快速验证端侧前处理逻辑不管转成什么格式最终部署时都有一个共识图像前处理必须和训练时一致。YOLOv5 训练时的前处理包括 letterbox 缩放、颜色通道转换 RGB、归一化到 0~1。很多人在端侧写前处理时忘了做 letterbox直接把原图 resize 到 640x640导致检测框位置偏移。验证方法很简单拿一张标注过的图跑完推理后把检测框画出来和原标注对比如果框整体偏移八成是 letterbox 没做或参数不对。import cv2 def letterbox(img, new_shape(640, 640), color(114, 114, 114)): shape img.shape[:2] r min(new_shape[0] / shape[0], new_shape[1] / shape[1]) new_unpad int(round(shape[1] * r)), int(round(shape[0] * r)) dw, dh new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] dw, dh dw // 2, dh // 2 img cv2.resize(img, new_unpad, interpolationcv2.INTER_LINEAR) top, bottom dh, dh (new_shape[0] - new_unpad[1] - dh) left, right dw, dw (new_shape[1] - new_unpad[0] - dw) img cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, valuecolor) return img这段 letterbox 代码和 YOLOv5 源码里的保持一致注意它先按比例缩放再在两边补边而不是直接拉伸。补边的颜色默认是灰色 114如果训练时改了颜色这里也要跟着改。从那以后我每次做端侧部署第一件事就是拿训练集里的一张图跑端侧推理打印出检测框坐标和 PyTorch 跑出来的结果做对比误差在几个像素内才算通过。这一步能挡住绝大多数前处理和后处理的低级错误希望帮到你。本文还有配套的精品资源点击获取