☰
YOLOv5模型转ONNX实战:部署加速与踩坑指南
2026/10/20 17:05:55 网站建设 项目流程

1. 为什么训练好好的模型一上生产就卡壳

1.1 PyTorch模型部署的三大痛点

如果你用YOLOv5跑过几个项目,大概率会经历这样的场景:训练时一切顺利,loss曲线漂亮,验证集的mAP也很理想。结果模型一交到部署环境,问题就来了——Python环境版本对不上,依赖库冲突,推理速度慢到无法接受,甚至换一台机器就报错跑不起来。

这不是某个人的技术失误,而是PyTorch模型部署本身的通病。

PyTorch模型自带两个明显的短板。第一是它对Python运行时的强依赖,模型权重文件(.pt)本质上需要加载进PyTorch的运行时才能执行推理,这意味着目标机器上必须安装对应版本的PyTorch、NumPy、OpenCV等一堆依赖。第二是它的性能调度开销,PyTorch的动态图机制在训练阶段是优势,但在推理阶段反而是负担,每次前向传播都要经过Python层的调用,这种开销在追求毫秒级延迟的场景下非常致命。

所以才需要引入ONNX(Open Neural Network Exchange)这类中间表示格式。ONNX做的事情很简单:把PyTorch训练好的模型权重和网络结构,转换成一个与框架无关的计算图描述文件。这个文件不依赖Python运行时,可以被各种推理引擎直接加载和执行,而且这些推理引擎普遍做了算子融合、内存复用、精度校准等优化,推理效率明显优于直接跑PyTorch。

1.2 ONNX作为中间格式的取舍逻辑

ONNX在工程界被称为“模型交换的通用语言”,这个说法虽然有点夸张,但它确实解决了模型从训练框架流向推理框架的最后一段路的问题。

以YOLOv5为例,转换后的ONNX模型可以交给ONNX Runtime处理,这是微软的开源推理引擎,CPU和GPU上的表现都不错;也可以转换成OpenVINO格式跑在Intel的硬件上;还能转换成TensorRT的engine文件跑在NVIDIA显卡上,实现FP16甚至INT8量化推理。也就是说,ONNX像是一个标准化的中转站,让同一个训练好的模型适配多种硬件和推理后端,而不用针对每个目标平台重新训练或重写推理逻辑。

这里要重点强调一个思路:ONNX不是终点,而是中间环节。

很多新手误以为转成ONNX就万事大吉,直接把ONNX文件拿去做生产推理。这种做法不是不行,但通常不会是最优选择。在实际项目中,ONNX更多扮演的是“可移植的中间格式”,真正的生产环境往往还会再把它转成目标硬件专属的高性能格式,比如TensorRT的engine、OpenVINO的IR文件、Core ML的mlmodel等。在本篇文章里,我会聚焦在PyTorch模型转ONNX这一步,但心里要有这根弦:ONNX只是桥,不是岸。

2. 转换前的环境准备与版本匹配

2.1 版本组合的黄金搭配

YOLOv5模型转换这件事,最容易被忽略但影响最大的因素就是版本匹配。网上能搜到大量转换报错的求助帖,粗略统计下来,超过一半的问题都源于torch和onnx的版本不兼容。

我实测下来比较稳妥的版本组合是这样的:

组件推荐版本备注
Python3.8 ~ 3.103.7以下对新版torch支持差
PyTorch1.10 ~ 1.132.x也可以,但个别算子导出需额外处理
ONNX1.12 ~ 1.16太老的版本缺算子支持
onnxruntime1.13 ~ 1.18用于转换后的验证推理
onnx-simplifier0.4.x简化计算图,减小模型体积

这里特别说明一下为什么PyTorch版本不建议使用太新的2.x。PyTorch 2.0引入了TorchInductor等编译机制,某些算子在trace导出时生成的ONNX计算图会变得非常臃肿,掺杂了大量辅助节点。YOLOv5官方仓库在很长一段时间里也只对1.x系列的torch做完整测试。如果你用的是新版本且转换时遇到算子兼容问题,最省事的解决办法不是硬啃报错,而是直接建一个Python 3.9 + torch 1.13的虚拟环境来做转换。

2.2 安装依赖需要避开的坑

依赖安装看起来是小事,但这里有一个经常踩的坑:torch和torchvision版本必须配套。

YOLOv5推理代码中会用到torchvision的ops模块,特别是nms(非极大值抑制)等相关操作。如果torch和torchvision版本对不上,轻则导入报错,重则在导出模型时出现莫名其妙的Segmentation Fault。建议使用官方提供的版本对照表来安装,比如torch 1.13.0对应torchvision 0.14.0。

另外还需要注意onnxruntime的安装。如果只是做验证,安装CPU版本的onnxruntime就足够了:

pip install onnxruntime

如果目标环境有NVIDIA GPU,可以换装onnxruntime-gpu。但做转换验证阶段,CPU版本反而更省心,原因是GPU版本还会牵扯CUDA、cuDNN的版本匹配,属于额外的变量。先把CPU版本跑通,后续再针对性的装GPU版本,这个顺序更符合排错逻辑。

安装完成后可以用一段简单的代码验证环境是否就绪:

import torch import onnx import onnxruntime as ort print("torch version:", torch.__version__) print("onnx version:", onnx.__version__) print("onnxruntime version:", ort.__version__)

如果这三个库都能正常导入且打印出版本号,说明基础环境没问题,可以进入下一步。

3. 核心转换步骤与关键参数逐个拆解

3.1 从yolov5s.pt出发的最简转换

YOLOv5官方仓库里自带了一个导出脚本export.py,在正常情况下,一条命令就能完成.pt到.onnx的转换:

python export.py --weights yolov5s.pt --include onnx

这条命令默认会做这些事情:加载yolov5s.pt权重文件,构建YOLOv5s网络结构,使用一组虚拟输入执行一次前向传播,通过torch.jit.trace记录计算图,再通过torch.onnx.export把计算图输出为ONNX格式。

但实际项目里,很少有人直接裸用这条命令。因为你训练出来的模型几乎都是自定义的类别数、自定义的输入尺寸,这些都需要通过参数明确指定。一个更贴近实际项目的导出命令是:

python export.py --weights runs/train/exp/weights/best.pt \ --img 640 \ --batch 1 \ --include onnx \ --opset 12 \ --simplify

各个参数的含义我用表格列一下,方便你对照检查:

参数作用不设置时的默认值推荐设置
hue类别数校准使用pt文件内嵌信息训练时保持对应
--img输入图片尺寸640与训练尺寸一致
--batch导出时的batch大小1实际部署时确认
--opsetONNX算子集版本1212或13
--simplify计算图简化不启用推荐开启

看起来很简单对吧?但真正的门道藏在细节里。

3.2 动态尺寸导出的实际意义

默认情况下,--img 640导出的模型只接受640×640的输入,一旦输入变成416×416或1280×1280,模型就会直接报错。这在批处理大小固定的场景下没问题,但如果是做流式视频检测,或者需要面对不同分辨率的图片,就会很难受。

要解决这个问题,可以开启动态输入尺寸:

python export.py --weights best.pt \ --img 640 \ --batch 1 \ --include onnx \ --opset 12 \ --simplify \ --dynamic

开启--dynamic后,ONNX模型对输入的分辨率就不再有硬性约束,可以接受任意H×W的输入。但这里必须说一句:动态尺寸不是免费的午餐。

ONNX Runtime在动态尺寸模式下,每次遇到新的输入分辨率都可能触发内部的内存重分配和算子形状推断,导致推理速度出现明显波动。我在实际测试中发现,固定尺寸的模型在640×640输入下CPU推理大约需要80毫秒,而动态尺寸模型在相同输入下可能要95毫秒以上,偶尔还会因为形状不匹配触发重新编译导致单次延迟飙到200毫秒以上。

所以我的建议是:如果你的部署场景对延迟敏感,尽可能使用固定尺寸;如果无法避免多分辨率输入,提前在代码里做好resize预处理,统一把输入控制在少数几个尺寸档位上,尽量别真的把动态尺寸完全放开。

3.3 batch维度需要注意的事项

batch这个参数在导出时容易被忽视,因为很多人潜意识里觉得batch=1就够用了。这个判断在大多数场景下是对的,但有几个特殊情况需要单独考虑:

  • 批量离线推理:如果部署环境需要一次性处理多张图片,batch=4或batch=8的模型会更高效。
  • TensorRT转换:后面如果打算把ONNX转到TensorRT,engine文件默认会绑定一个固定的batch size。建议在导出ONNX时就确定好最终要用的batch值。
  • 动态batch:ONNX也支持batch维度动态化,即允许同一个模型接受不同的batch输入。用法是在--dynamic参数后面追加batch信息,但跟动态尺寸一样会带来额外的运行时开销。

就个人经验而言,绝大多数中小项目用batch=1就够了。ONNX Runtime本身支持对batch=1的模型做比较高效的调度优化,没必要为了“以后可能能用上”多线程而把batch设大,毕竟这会显著增加显存或内存占用。

4. 转换后的验证与踩坑排错链路

4.1 输出结构校验不能只看“有没有生成文件”

很多同学转完模型后,看到目录下多了一个.onnx文件就认为大功告成。这个习惯非常危险。ONNX文件能生成,只能说明torch.onnx.export执行完了,不代计算图的结构一定正确。

正确的校验方式分两步:

第一步,使用onnx库检查模型结构是否合法:

import onnx model = onnx.load("best.onnx") onnx.checker.check_model(model) print("model is valid")

如果这一步通过,说明ONNX文件本身没有语法或结构上的错误。但这还不够,因为它只能验证“静态图”的正确性,无法保证推理结果和原始PyTorch模型一致。

第二步,用onnxruntime加载模型,对比输出差异。

import numpy as np import onnxruntime as ort import torch # 随机生成一张与训练尺寸一致的输入 dummy_input = np.random.randn(1, 3, 640, 640).astype(np.float32) # 用ONNX Runtime推理 ort_session = ort.InferenceSession("best.onnx") ort_outputs = ort_session.run(None, {"images": dummy_input}) # 用原始PyTorch模型推理 model = torch.load("best.pt", map_location="cpu")["model"].float().eval() with torch.no_grad(): torch_outputs = model(torch.from_numpy(dummy_input)) print("ONNX output shape:", [o.shape for o in ort_outputs]) print("PyTorch output shape:", [o.shape for o in torch_outputs])

这一步的作用是确认模型转换后没有丢失精度或改变输出行为。正常情况下,两种推理方式得到的输出shape应该完全一致,数值上会有极其微小的差异(通常在1e-5量级),这是浮点数运算顺序不同导致的正常误差,不影响检测结果。

4.2 输出张量的shape解读与协议约定

YOLOv5的ONNX模型输出是什么样子的?这里需要特别说明一下,因为很多人在拿到输出后不知道怎么解析。

以YOLOv5s为例,输入为1×3×640×640时,ONNX模型的输出shape通常是1×25200×85。这个25200是三个检测头各自输出的anchor数量总和:(80×80 + 40×40 + 20×20) × 3 = 25200。85这个数字则是由80个类别概率 + 4个边界框坐标 + 1个目标置信度组成的。

但这里有个非常重要的坑:ONNX输出格式和PyTorch原始输出的排列顺序可能不同。YOLOv5的PyTorch模型在推理时,内部会把检测结果整理成[x_center, y_center, width, height, obj_conf, class_scores...]的格式,而导出ONNX时需要有一个后处理步骤把坐标从相对网格坐标转换成相对于原图的像素坐标。YOLOv5官方仓库的export.py会自动完成这个pre-processing和后处理的打包,所以导出的ONNX模型输出的坐标是已经换算过的像素坐标。

如果你是自己手写导出脚本,而不是用官方export.py,就很容易在这一步栽跟头。我记得有人用自写脚本导出后,在ONNX Runtime里拿到的坐标全部是0到1之间的比值,直接拿去做可视化时发现所有框都画在了图片左上角。排查到最后才发现,是导出时没有把decode操作包含进计算图导致的。

4.3 常见报错与解决办法

转换过程中有几个高频报错,我把它们列出来,并附上我实际排查时的解决思路。

报错一:Unsupported operator或者Constant node not in opset

这个报错的核心原因是模型中的某些操作在当前ONNX算子集版本里没有对应实现。解决办法通常有两种:要么提高--opset版本,比如从12改成13或14;要么检查模型中是否有自定义操作,需要将自定义操作手动替换为ONNX支持的等效操作。

报错二:Exported graph contains cycles

这个在YOLOv5里不太常见,但如果出现了,大概率是因为模型中存在循环依赖或控制流操作,torch.jit.trace无法正确处理。最简单的解决思路是升级到最新版PyTorch和ONNX,官方迭代过程中修复了很多这类边际案例。

报错三:Segmentation Fault during export

这类崩溃多见于torch和torchvision版本不匹配,或者CUDA版本与PyTorch版本冲突。排查方法是先跑一遍官方自带的检测demo,确认PyTorch本身能正常工作,再导出;如果PyTorch推理正常但导出崩溃,试着把模型切到CPU模式再导:

model.model.eval() model.model.fuse() model.model = model.model.to("cpu")

这种办法很多时候能绕过CUDA相关的导出路径问题。

4.4 用onnx-simplifier进一步压缩模型

转换完成后,我强烈建议运行一下onnx-simplifier。它的作用是对计算图做冗余节点消除、算子融合、常量折叠等优化,通常能让ONNX文件的体积缩小20%到30%,同时推理速度也有轻微提升。

安装很简单:

pip install onnx-simplifier

调用方式:

python -m onnxsim best.onnx best_simplified.onnx

简化前后的模型我都做过对比测试,在相同硬件上的推理延迟差异在几个百分点以内。但它真正的意义在于“净化”计算图,方便后续转TensorRT或OpenVINO时减少兼容性问题。很多推理框架自己的优化器在遇到复杂图结构时会报错,简化后的模型则能顺利通过。

5. 从ONNX到实际部署的衔接

5.1 推理时的NMS到底该放在模型的哪一侧

这是YOLOv5转ONNX后大家问得最多的问题之一:ONNX模型里到底要不要包含NMS操作?

为了说清楚这个问题,先明确NMS(非极大值抑制)在目标检测中的作用。YOLOv5的三个检测头会输出25200个候选框,其中绝大多数是重复检测的冗余框。NMS的作用是扫描所有框,根据IoU和置信度阈值保留最优框、剔除冗余框。

YOLOv5官方仓库提供的export.py默认不把NMS放进ONNX模型里。也就是说,导出的ONNX模型只负责输出“原始的25200个框的置信度和坐标”,NMS需要由调用方在外部完成。这样做的好处是模型结构干净,ONNX导出和推理框架转换都更稳定;坏处是部署方需要自己实现NMS逻辑,或者在ONNX Runtime的后处理里额外做一层处理。

与之相对的另一种做法是使用torchvision.ops.nms或自定义算子把NMS直接集成进模型,在导出时一并打包。这种方式可以让ONNX模型“自包含”,但不同推理框架对自定义NMS算子的支持程度不一样,转TensorRT时很容易翻车。

我的建议是:如果走TensorRT路线,ONNX里绝对不要包含NMS,后续在TensorRT里使用EfficientNMS插件来实现;如果走ONNX Runtime路线且对性能要求没那么苛刻,可以把NMS放在外部后处理中,用NumPy实现,容易调试也容易替换。

5.2 不同推理框架的适配差异

ONNX转换完成后,常见的目标框架有几个,适配要点各不相同:

目标框架主要适配动作注意事项
ONNX Runtime直接用性能较好,CPU/GPU均可
TensorRT再转engine推荐固定尺寸和batch
OpenVINO再转IR格式Intel CPU上优势明显
Core ML再转mlmodelApple Silicon部署
NCNN再转ncnn格式移动端/嵌入式场景

以TensorRT为例,转换为engine文件时经常遇到两类问题:一是某些ONNX算子TensorRT不支持,如Upsample的特定模式;二是计算图中有动态shape导致engine编译失败。这时候在前面做了onnx-simplify的好处就会体现出来,简化后的模型兼容性明显更好,编译失败的概率大幅降低。

5.3 一些性能观察与优化思路

在文章最后分享一些我实测的优化想法。

同样是YOLOv5s权重,用PyTorch直接推理,CPU上单张640×640图片大约需要350毫秒;转成ONNX后用ONNX Runtime CPU推理,大约能压到80毫秒左右;如果进一步转成TensorRT FP16,在桌面级GPU上可以做到3到5毫秒。这个差异足以说明ONNX转换本身并不是多此一举,而是整个部署流程中收益最明显的一步。

另一个常见的优化手段是用FP16精度导出。YOLOv5的export.py支持--half参数,导出的ONNX模型在支持FP16的硬件上推理速度能再提升一倍左右。但要特别注意:FP16模型在CPU上无法直接运行,而且某些自定义算子对FP16的支持不完善,导出后需要用onnxruntime-gpu实测一遍再确认。

还有关于opset版本的选择,我建议默认用12或13,这两个版本在兼容性和算子支持之间平衡得最好。opset 16以上的版本虽然多了不少现代算子,但在转TensorRT和OpenVINO时踩坑的概率明显更高,除非你的模型中真的用了需要新算子集才能表达的语法,否则没必要追新。

转换这一步做到位了,后面的部署之路就会顺畅很多。我在多个项目中用过同一套流程——固定尺寸导出、ONNX Runtime验证、onnx-simplify简化、再按目标框架转换,这套路径的风险最小,出了问题也最容易定位。希望对正在捣鼓YOLOv5部署的你有所帮助。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询