做了一段时间目标检测部署我最大的感受是训练和部署完全是两个世界。训练时PyTorch模型怎么方便怎么来一放到部署环节就得面对各种推理引擎而ONNX往往是绕不开的中间层。YOLOv5算是部署生态比较完善的检测模型了很多朋友第一次接触ONNX就是想把训练好的.pt权重转成ONNX文件再交给后续推理框架使用。这类转换如果只看官方文档好像执行一行python export.py --weights best.pt --include onnx就完事了但真正动手实操过的人都知道这里面的门槛远不止一条命令。这篇内容我就把YOLOv5从PyTorch模型转换为ONNX文件的完整过程、核心思路和踩过的坑一次说清楚适合刚接触模型部署的开发者也适合已经转完但推理结果对不上的朋友对照排查。1. 为什么要把YOLOv5的PyTorch模型转成ONNX1.1 ONNX在部署链路上到底是什么角色ONNX全称是Open Neural Network Exchange翻译过来就是开放神经网络交换格式。你可以把它理解为深度学习框架之间的“通用插座”PyTorch、TensorFlow这些框架各有各的脾气手里的模型格式也互不相同但ONNX定义了一套统一的计算图描述规范把网络结构、权重、算子和输入输出定义统统固化在一个文件里。得益于这套统一规范ONNX成了模型从训练框架走向硬件推理引擎的必经中转站。在YOLOv5的实际部署场景里ONNX的典型路径是这样的首先把训练好的.pt权重导出成.onnx文件然后根据目标硬件选择不同的后续转换工具。比如英伟达显卡上用TensorRT加速Intel CPU上用OpenVINO通用场景用ONNX Runtime。我实际测试下来的结论是从ONNX转TensorRT和OpenVINO都比较顺畅路径也成熟。所以很多项目干脆把ONNX作为模型分发格式只要交出一个.onnx文件下游怎么用都能接得住。1.2 PyTorch模型直接部署的痛点在哪里有人可能会问既然PyTorch模型也能加载推理为什么非要转一道直接部署PyTorch模型确实简单torch.load加model.eval()就能跑。但真实生产环境里这种做法会带来一连串麻烦。第一PyTorch运行时太重了。完整推理依赖整个torch库动辄几百MB放到边缘设备或者嵌入式设备上根本不现实。第二推理效率不理想。PyTorch默认的推理路径是为了训练灵活性设计的带了很多自动求导和动态图相关的额外逻辑虽然model.eval()能关掉一部分但计算图依然不够“纯净”很难被底层引擎做深度优化。第三跨语言部署困难。PyTorch推理几乎绕不开Python环境而实际生产线上的服务端可能用C、Java甚至Go总不能让所有服务都塞一个Python进程。ONNX模型本质上是一个静态计算图结构清晰、权重内嵌、不依赖训练框架天然就适合被推理引擎加载和优化。这也是为什么几乎所有硬件厂商都愿意优先适配ONNX导入。1.3 YOLOv5官方为什么把ONNX当作第一站YOLOv5官方提供的导出脚本里支持导出的格式非常丰富TorchScript、ONNX、TensorRT、OpenVINO、CoreML、TF.js等。但这么多格式里ONNX始终是承上启下的核心环节。原因有三点。一是生态覆盖极广。只要是有点名气的推理引擎基本都支持ONNX导入一次导出可以适配多个硬件平台。二是PyTorch转ONNX的工具链相对成熟。torch.onnx.export这套机制经过多年迭代对常见算子的兼容性已经很稳定YOLOv5的网络结构也不复杂导出成功率和还原度都很高。三是后续优化的想象空间大。ONNX模型可以继续做量化、剪枝、算子融合转成更适合特定硬件的格式。我自己的习惯是先导出ONNX做一次数值一致性验证确认模型没转坏再决定下一步走TensorRT还是其他方案这样能提前拦截很多部署问题。2. 环境准备与依赖安装2.1 建议的版本组合模型转换这件事版本搭配不合理是真的会卡死人的。我为这个坑掉过好几次这里直接给出我实测比较稳的组合省得你从头试错。PyTorch建议核心版本在1.13左右不要用太老的版本。1.8以下对ONNX导出的算子映射支持不太全容易出现莫名其妙的“算子不存在”报错。PyTorch 2.0及以上也能导出YOLOv5但要注意新版算子集中引入了一些新op比如scaled_dot_product_attention对onnx版本要求更高环境如果没跟上导出时会直接报错。onnx库建议至少1.12.0新版更好。onnx-simplifier用于简化模型我后面会细说。验证推理结果需要onnxruntime或onnxruntime-gpuGPU版本在英伟达显卡机器上安装。如果你的设备没有GPU就装CPU版本。Python环境用3.8到3.10都没问题太新的Python版本反而容易和onnxruntime的部分旧版本出现兼容问题。再强调一点如果用的是PyTorch 2.x优先把onnx升级到较新版本不要抱着旧环境不动。版本适配是做部署的基本功课省得在后面排查时间上花大成本。2.2 获取代码库与依赖YOLOv5的官方代码仓库直接用git clone拉下来就行。这里建议你关注版本分支不要随意用最新主线因为不同版本的export.py参数定义和默认行为可能不一样。我自己长期用v6.0和v7.0这两个分支导出逻辑都比较稳定。拉取代码后进入项目根目录安装依赖。我一般这样做pip install -r requirements.txt pip install onnx onnx-simplifier onnxruntimerequirements.txt里已经包含了PyTorch和基础依赖但ONNX相关库通常不在默认列表里需要单独装。如果你需要GPU推理验证就把onnxruntime换成onnxruntime-gpu注意两者不能同时混装否则会有DLL冲突之类的报错。安装完成后最好确认关键库的版本号防止后面排查问题时搞不清环境状态。python -c import torch, onnx, onnxruntime; print(torch.__version__, onnx.__version__, onnxruntime.__version__)2.3 先跑通一条最小导出链路我建议新到手的部署环境不要上来就导自己的模型先拿官方预训练权重做一遍最小链路验证。具体做法是把yolov5s.pt下载下来然后执行一条最简单的导出命令python export.py --weights yolov5s.pt --include onnx --img-size 640 640如果这一步顺利通过说明环境基本没问题。如果这一步都报错那大概率是版本兼容性问题先解决环境再继续。千万别跳过这个步骤。很多朋友直接拿自己的best.pt导出失败查了一整天最后发现是torch和onnx版本不匹配换版本后一条命令就过了。3. YOLOv5官方导出脚本的完整实操3.1 export.py核心参数逐个说清楚熟悉export.py的参数是做转换的基本功。我挑几个最常用的参数说明一下容易忽略的细节也会标注清楚。--weights指定待转换的权重文件路径一般指向runs/train/exp/weights/best.pt。--include指定导出格式用onnx就行脚本会自动识别并调用对应的导出逻辑。--img-size是输入图像尺寸默认是640 640大部分场景都不用改。如果你训练时用960分辨率这里就得对应改成960 960。--batch-size指定导出模型的批大小YOLOv5导出时默认batch1。这里要注意批大小会影响导出的输出形状后续部署时一般也是单张推理生产环境里用batch1就够了。--dynamic控制是否导出动态维度开启后模型输入允许不同分辨率后面我会专门展开讲动态和静态的区别。--simplify会调用onnx-simplifier对模型做简化可以移除部分冗余算子让模型更干净。--opset指定ONNX算子集版本YOLOv5默认是12建议用12或13。--half导出半精度FP16模型如果你确定部署端支持FP16再开否则可能给自己挖坑。--device指定导出时用的计算设备cpu最稳0用GPU。3.2 最常用的导出命令与分析我用得最多的导出命令是这一条python export.py --weights runs/train/exp/weights/best.pt --include onnx --img-size 640 640 --opset 12 --simplify --dynamic一个个拆开解释。--weights指向自己训练好的权重--include onnx锁定输出格式--img-size 640 640固定看训练时的分辨率--opset 12保持合理的算子集版本适配性最好--simplify开启模型简化--dynamic开启动态尺寸支持。这样导出的模型在部署时允许输入不同尺寸的图像不需要因为分辨率变化重新导出。不过要注意动态模型对部分推理引擎不够友好比如TensorRT虽然支持动态shape但需要显式配置优化范围。如果你后续确定只用固定尺寸比如统一640x640输入那--dynamic完全可以不加静态模型在不同引擎上的兼容性更好。如果转换过程中报错先尝试去掉--simplify再不行就去掉--dynamic。把变量减到最少能更快定位是哪一步出的问题。3.3 转换后的模型文件是什么样的转换完成后在权重同目录下会多出一个.onnx文件比如best.onnx。这个文件的大小通常比原始的.pt小一些原因是.pt文件里除了模型结构还包含优化器状态、训练轮次、类别名等大量与推理无关的信息而.onnx文件只保留计算图和必要的权重参数。可以用ONNX库读取模型文件快速确认输入输出定义是否正常import onnx model onnx.load(best.onnx) onnx.checker.check_model(model) print(model.graph.input) print(model.graph.output)这两段代码会打印模型的输入节点和输出节点。正常情况输入名一般是images输出名一般是output。如果图结构有问题check_model会直接抛异常这是第一时间发现模型损坏的检查手段。4. 转换中的核心原理与关键细节4.1 静态尺寸和动态尺寸怎么选这是模型转换里最容易被忽略、也最容易影响部署结果的选择。静态尺寸指的是导出时把输入shape固定下来比如固定为(1, 3, 640, 640)。动态尺寸则允许在推理时传入不同的宽高比如同一份模型既能跑640分辨率又能跑1280分辨率。动态尺寸的实现原理并不复杂。torch.onnx.export在导出时通过dynamic_axes参数声明哪些维度是可变的比如dynamic_axes{images: {0: batch, 2: height, 3: width}}这样导出的模型输入shape里的batch、height、width都变成了符号变量推理引擎可以接受任意合法值。但动态尺寸不是免费的。首先输出层的shape也会跟着变成动态这会直接影响后处理代码的编写。其次部分推理引擎对动态shape的支持不如静态shape成熟可能引入额外配置甚至性能倒退。所以我通常建议检测场景输入分辨率基本固定的用静态尺寸需要处理多尺度输入的用动态尺寸。不要为了留后路盲目开动态后期有的是坑在等你。4.2 Batch维度怎么处理Batch维度在目标检测部署里一般不是关注重点因为线上推理几乎都是单张图片走一遍batch1够用。导出时通过--batch-size 1固定即可。但如果导出时用batch1而推理引擎配置的输入batch是4就会直接报维度不匹配。反过来导出时用batch4单张推理时虽然可以复用但白白浪费了计算资源。我建议单卡部署就无脑batch1。特殊场景比如同时处理多路视频流可以考虑把多帧拼成batch去推理但那是另一套优化思路了和ONNX转换关系不大。4.3 输出层的形状和后处理的关系YOLOv5在640x640输入下输出层经过Detect头处理后默认输出一个拼接好的检测结果张量常见形状是(1, 25200, 85)或者(1, 3, 8400, 85)具体取决于你用的版本。其中25200或8400代表所有尺度的候选框总数85代表4个坐标信息加1个目标置信度再加80个类别得分。坐标信息在输出时通常是xywh格式也就是中心点x、中心点y、宽度、高度。类别得分需要经过sigmoid激活后才能变成真正的概率值。转换ONNX时Detect头内的后处理逻辑一般不会全部合入ONNX坐标解码和NMS往往保留在外部完成。原因很简单NMS的动态性太强很难用静态计算图高效表达而且不同的部署框架对NMS的优化方案也不一样。理解了输出层的形状和含义你在用ONNX推理后写后处理时就不会一头雾水。4.4 简化模型到底有没有必要--simplify的作用是把ONNX模型里的常量计算提前折叠、把一些冗余算子合并整体看起来更简洁推理引擎加载速度也会快一些。这个步骤听起来只有好处没有坏处但我实际使用中遇到过一次意外某些版本的onnx-simplifier对动态shape模型简化后输出结果出现了一点偏差。原因是要把动态维度抽象成符号计算简化器对这类图结构的处理偶尔会引入误差。所以我的建议是优先做简化但简化完之后一定要做数值一致性验证。如果发现简化后的模型输出和PyTorch原模型输出差异较大就放弃简化直接导出原始ONNX。多数情况下简化后的模型没有问题但这一步验证不能省。5. 常见问题与排查技巧实录5.1 导出后推理结果对不上怎么办这个问题是出现频率最高的。模型转换成功了ONNX Runtime也跑通了但检测框和原PyTorch模型的结果对不上框的位置、置信度都有差异。我遇到这种情况第一反应不是怀疑导出的计算图有问题而是先检查数据预处理是否一致。YOLOv5训练时一般会把图像缩放到640x640并将像素值除以255归一化到0到1之间。用ONNX推理时同样必须做完全相同的缩放和归一化。很多“结果不对”的案例其实是推理前把图片用cv2.imread读进来后没有除以255而PyTorch的DataLoader暗地里做了归一化两边不统一输出自然对不上。另外需要确认推理时的letterbox方式是否一致。YOLOv5默认会做等比例缩放并填充灰色边推理端如果直接用resize拉伸目标几何关系就已经变了。检查这类问题最简单的办法是我下面要说的数值一致性验证。5.2 数值一致性验证怎么做我每次导出ONNX后都会先跑一遍纯数值对比。方法很简单用同一个随机输入分别喂给PyTorch模型和ONNX Runtime比对输出张量的最大绝对误差。import torch import onnxruntime as ort import numpy as np # 原始PyTorch模型 pt_model torch.load(best.pt, map_locationcpu)[model].float() pt_model.eval() img torch.randn(1, 3, 640, 640) with torch.no_grad(): torch_out pt_model(img)[0].numpy() # ONNX Runtime推理 sess ort.InferenceSession(best.onnx, providers[CPUExecutionProvider]) ort_out sess.run([output], {images: img.numpy()})[0] diff np.max(np.abs(torch_out - ort_out)) print(fmax diff: {diff:.6f})如果max diff在1e-4量级以内基本可以认为转换成功。如果量级达到1e-2甚至更大说明转换过程中出现了实质性的计算差异。先用随机输入做对比好处是排除了图片预处理和后处理的干扰纯粹对比计算图本身。如果随机输入都过不了就不用在真实图片上浪费时间了。5.3 算子兼容性报错导出ONNX时经常会看到这样几类报错Op not registered、Unsupported opset version、Failed to export an ONNX attribute。这些问题的根源大多是版本不匹配。举个例子PyTorch 2.0导出的模型可能包含新引入的算子而老版本onnx无法解析此时需要升级onnx或者降低PyTorch版本。还有一类情况是opset设置太高而部署环境使用的推理引擎不支持那个版本的算子集。比如TensorRT某些旧版本只支持ONNX opset 13你导出时用了opset 17结果导入TensorRT就报错。解决办法也很直接把--opset降到12或13重新导出。遇到这类报错排查顺序是先确认torch和onnx版本环境再把--simplify去掉试试最后再调整--opset。一层一层排除基本都能定位到原因。5.4 不通过官方脚本的手动导出方法有时候官方脚本满足不了需求。比如自己改了YOLOv5的网络结构或者Detect头部分被自定义逻辑替换了官方脚本的导出路径可能直接失效。这时候就需要手动用torch.onnx.export导出。手动导出需要注意几点。模型必须先设为eval模式并关闭梯度需要准备一个dummy_input大小要和训练时的输入一致导出前要确认模型的前向输出是什么结构。示例代码如下import torch from models.experimental import attempt_load model attempt_load(best.pt, map_locationcpu) model.eval() dummy_input torch.randn(1, 3, 640, 640) torch.onnx.export( model, dummy_input, custom.onnx, input_names[images], output_names[output], opset_version12, dynamic_axes{images: {0: batch, 2: height, 3: width}}, )这里attempt_load是YOLOv5代码库里的加载函数和普通的torch.load不同它会正确处理模型结构。如果自定义了后处理节点也可以把部分后处理逻辑一并封装进导出流程但前提是这些逻辑能被ONNX算子表达。类似NMS这种包含循环和动态数量的操作ONNX表达起来就非常费劲一般不建议塞进去。5.5 踩坑速查表这里把我在实际项目中遇到的典型问题整理成一张表方便你快速对照。现象可能原因解决办法导出时报算子不认识的错误onnx版本过低或torch版本过高升级onnx或降低torch版本导出后模型很大包含训练状态和优化器用官方export.py导出它会剔除训练相关参数ONNX Runtime报错不支持某种算子opset版本超出运行时支持范围将--opset降到12或13简化后推理结果偏了onnx-simplifier对动态shape处理有误差不使用--simplify重新导出并验证half导出后推理精度下降部署端FP16支持不完善改用FP32导出动态尺寸模型无法在某个引擎加载引擎要求静态输入重新导出固定尺寸模型5.6 一个不算常见但很隐蔽的坑还有一个小坑值得提一下。如果训练时修改了模型类别数但导出时误用了官方预训练模型的配置文件导出的ONNX输出维度会不对后处理时期望80类实际输出却只有原来模型定义的类别数。这类问题表面上看是维度不匹配实际上是你加载权重时对应的模型定义搞错了。排查方法是打印ONNX模型的输出shape和你训练时的类别数加5做对比。如果对不上就去检查data.yaml中的类别列表数量以及导出时加载的模型结构是否正确。这类问题不容易被发现但一旦发生会浪费大量时间。写在后面的一点心得转换ONNX这件事表面上只是一条命令的事但真正让我觉得有价值的是理解这条命令背后模型结构、算子映射和部署需求之间的博弈。我个人的习惯是每次训练出一个新模型先导出ONNX做数值验证确认无误后再去部署其他平台。这个习惯帮我挡掉了至少一半的诡异问题。最后再分享一个小技巧导出的ONNX文件用netron打开看一遍图结构能直观确认输入输出、算子流程和你预期是否一致。不要怕麻烦部署这件事前期多花十分钟检查后期能省一整天排查时间。