ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

YOLOv5模型导出实战:从pt到ONNX与TensorRT engine避坑指南

YOLOv5模型导出实战:从pt到ONNX与TensorRT engine避坑指南

训练跑完,best.pt躺在权重目录里,几百兆的文件看着挺唬人,但只要下游同学一句“给我一份 onnx,再给我一份 TensorRT 的 engine”,很多人第一次碰 export.py 就会原地卡住半天。我见过太多这样的情况:有人把命令行参数抄错一个字母,导出的 onnx 用 onnxruntime 一跑全是乱码框;有人把 engine 文件拷到另一台机器上,加载直接抛异常,然后怀疑是文件损坏;还有人导出成功了,但精度掉得莫名其妙,最后发现是输入尺寸没和部署端的预处理对齐。这篇就把 YOLOv5 仓库里自带的 export.py 这条路走通一遍,从 pt 到 onnx 再到 TensorRT 的 engine,把每一步“为什么这么设”“不这么设会怎样”讲清楚,重点放在那些文档里不会写、只有真踩过才知道的细节上。适合已经能跑通训练、接下来要往部署端交付的读者,也适合手上有一堆历史权重需要批量转格式的同学。

1. pt 到 onnx 再到 engine:先弄清 export.py 到底在干什么

1.1 它不是格式转换按钮,而是一条带前置条件的构建流水线

很多人对 export.py 的心理预期是“选个格式、点一下、出文件”,实际上它内部是一条有明确顺序和前置条件的流水线。脚本先按--weights加载权重并重建模型结构,然后把模型切到推理态、把检测头里的 anchor 网格预生成逻辑改成导出友好模式,接着根据--include后面的格式列表逐个调用对应的导出函数,每导出一个格式就把产物落到权重同目录,最后由一个汇总函数把文件复制到--project/--name指定的目录并打印一张结果表,列出成功、失败和文件大小。这一串动作里任何一环的条件不满足,后面全断。

所以要建立的第一条认知是:--include onnx engine写在一起的时候,这两个格式是串行独立执行的,并不是先出 onnx 再由 onnx 编译 engine。也就是说 onnx 导出失败不影响 engine 导出,engine 导出失败也不会把已经生成的 onnx 删掉。这一点在实际排错时特别有用——当你看到终端里 onnx 那行打了勾、engine 那行报错,就不用回头怀疑权重文件本身有问题,问题一定出在 TensorRT 相关的环节。

1.2 导出的图里没有 NMS,这是最常见的认知错位

必须提前说清楚的一件事:YOLOv5 导出的 onnx 和 engine,输出的是未经非极大值抑制的原始预测张量,形状大致是(batch, 候选框数量, 5 + 类别数)。候选框数量取决于输入尺寸,640×640 输入下是 25200。很多第一次做部署的人拿到这个输出,喂给可视化脚本发现框叠成一片,就以为是导出过程把权重搞坏了,其实完全不是——NMS 从来就没被包含在导出图里,它需要在推理端自己实现,或者用仓库推理脚本里现成的后处理。

这个设计不是偷懒,是有明确理由的。NMS 里含有大量动态循环和排序操作,把它塞进计算图会让 TensorRT 这类静态编译器非常难受,要么编译不出来,要么编译出来性能比不融合还差。把后处理留在宿主侧,计算图保持纯粹的卷积加激活,编译器才能做充分的层融合和内核选择。理解了这一点,你在验证导出结果时就不会拿“有没有框”当标准,而是去对齐原始张量的数值。

1.3 onnx 是图翻译,engine 是针对显卡的离线编译

这两个格式的定位差得非常远,混为一谈会吃大亏。onnx 是一份平台中立的计算图描述文件,内部记录的是算子序列、张量形状和权重常量,它本身不能执行,需要由运行时去解释。所以 onnx 的可移植性很好,同一份文件可以在不同机器、不同推理后端上跑,代价是运行时要做算子映射和调度,性能和显存占用都不是最优。

engine 则是 TensorRT 在当前这台机器上,针对当前这块显卡的计算能力、当前装的 TensorRT 版本,把计算图编译成的一份二进制执行计划。编译过程中它会做层融合、精度选择、内核自动调优、显存预分配,所以跑起来快、占用省,但代价是这份文件被焊死在了这套环境上。换一块不同架构的显卡、换一个 TensorRT 大版本,甚至换一个不兼容的驱动,engine 都可能反序列化失败。把 engine 当成“更快的 onnx”到处拷贝,是踩坑率最高的一个做法。

记住这条:onnx 可以进版本库、可以交付、可以跨机器传;engine 只属于构建它的那台机器,交付时要么给 onnx 让对方自己构建,要么把构建脚本一起打包。

2. 环境钉版本:这一步偷懒,后面全是玄学报错

2.1 CUDA、TensorRT、PyTorch 三者的配套关系

导出这条链路上最容易被忽视、又最容易出问题的就是版本配套。PyTorch 的 CUDA 运行时版本、系统装的 CUDA 工具链版本、TensorRT 编译时依赖的 CUDA 版本,这三者需要大致处在同一个大版本线上,否则典型症状是:onnx 能导出,engine 导出时在某个环节抛出一句非常笼统的异常,看不出根因。

我的做法是先固定 TensorRT 版本,再倒推其他。因为 TensorRT 的安装包本身就绑定了它所支持的 CUDA 大版本,你装哪个 wheel,CUDA 的可用范围就定了;然后选一个与这个 CUDA 版本匹配的 PyTorch 官方构建;最后确认 Python 版本在 TensorRT 的 wheel 支持列表里。这张对照关系不要凭记忆,直接去看 TensorRT 安装包附带的说明和你下载的那个 wheel 的文件名——文件名里通常就写着它对应的 Python 版本、CUDA 版本和操作系统平台,这是最可靠的信息源。

至于 YOLOv5 仓库本身,它对版本的要求其实不严,真正严格的只有 TensorRT 那一侧。但仓库的依赖文件里往往会带一个 onnx 和一个 onnxruntime 的版本下限,这些也要一并对齐,否则--simplify那一步可能因为简化器的接口变化而失败。

2.2 用独立虚拟环境把导出工具链隔离起来

导出环境强烈建议单独开一个虚拟环境,不要和训练环境混用。原因很实际:训练环境里装的 torch 版本、CUDA 版本是为了训练稳定选的,导出环境需要的是和 TensorRT 匹配的版本,两者经常对不上。混在一起的结果就是每次想切任务都得升降级,早晚把训练环境搞坏。

隔离之后还有一个好处,你可以把导出环境完整地写进一份依赖清单,别人拿到权重后按清单重建环境,导出的结果和你本地的完全一致。这份清单里至少要钉住:torch 版本、onnx 版本、onnxruntime 版本、onnx-simplifier 版本、TensorRT 版本、PyCUDA 版本。全部写死具体版本号,不要用>=这种范围写法,因为导出结果的可复现性完全依赖于这几个包的精确版本。

在实际项目里我甚至会把这个环境做成一个容器镜像,每次导出都从镜像起一个干净的实例。这样做的好处是“导出失败”这件事变成了一个可以被复现的确定性问题,而不是“我本地是好的,你那边不知道怎么回事”。

2.3 三条自检命令,确认环境真的可用

环境装完别急着导模型,先用三条命令确认基础能力都在。这三条分别验证 PyTorch 能否看到显卡、TensorRT 能否正常初始化和打印版本、PyCUDA 能否建立 GPU 上下文。前两条不过就是环境没装好,第三条不过则是一个非常典型的问题——TensorRT 的 Python 绑定在部分导出路径上需要 PyCUDA 做显存指针绑定,缺了它引擎导出会在中途失败。

import torch print("torch:", torch.__version__) print("cuda available:", torch.cuda.is_available()) print("cuda version:", torch.version.cuda) print("device:", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "none") import tensorrt as trt print("tensorrt:", trt.__version__) logger = trt.Logger(trt.Logger.WARNING) builder = trt.Builder(logger) print("builder ok:", builder is not None) import pycuda.autoinit import pycuda.driver as cuda print("pycuda ok, device count:", cuda.Device.count())

三条都通过再往下走。这里有个细节值得提醒:import pycuda.autoinit会在导入时立刻初始化 GPU 上下文,如果你的脚本后面还要用 PyTorch 分配显存,两套上下文并存时偶尔会互相干扰。所以自检脚本跑完就退出,不要把它和真正的导出逻辑写在一个文件里。

3. 导出 ONNX:参数不是抄的,是算出来的

3.1 一条完整命令与每个参数的实际作用

先把一条能直接用的命令摆出来,然后逐项解释。这里用最小模型举例,实际换成你自己的权重路径即可。

python export.py \ --weights runs/train/exp/weights/best.pt \ --include onnx \ --imgsz 640 640 \ --batch 1 \ --opset 12 \ --simplify \ --device 0 \ --project runs/export \ --name onnx_fp32_b1

--weights是权重路径,支持.pt也支持直接传.onnx做二次处理。--include后面跟格式名,多个格式用空格分隔,写成--include onnx engine就是两个都导。--imgsz指定推理分辨率,可以给一个数(正方形)也可以给两个数(高宽)。--batch是导出时固化的批大小。--opset是 onnx 算子集版本。--simplify调用 onnx-simplifier 做图简化。--device指定用哪块显卡,写0表示第一块,写cpu则纯 CPU 导出。--project和--name控制产物目录。

不同小版本的参数名有细微差异,比如有些版本里分辨率参数叫--img或--img-size,批大小叫--batch-size。我的习惯是每次导出前先跑一遍python export.py --help把参数表打出来看一眼,这个动作只要五秒钟,但能省掉因为参数名不识别而导致的静默失败——有些版本对未知参数不会报错,而是直接忽略,你以为设置了其实没生效。

3.2 imgsz 和 batch 怎么定:跟部署端的输入源对齐

分辨率不能拍脑袋定。有两个硬约束:第一,必须是 32 的整数倍,因为主干网络做了五次下采样,最大步长是 32,非 32 倍数的尺寸会在特征图对齐时出问题;第二,要和部署端实际的输入预处理完全一致,包括 letterbox 的填充方式。

我在一个项目里遇到过精度掉得莫名其妙的情况,查了很久最后发现:训练和导出用的是 640×640,部署端为了省算力把输入改成了 512×512 但没有重新导出,而是让推理框架自己做动态尺寸适配。结果模型的 anchor 网格是按 640 预生成的,尺寸一变,网格和特征图对不上,框位置整体偏移。分辨率一旦确定就不要在推理端偷改,改了就重新导一份。

batch 的选择逻辑更简单:部署端一次处理几张图,就导出几。如果部署端是视频流逐帧处理,导 batch=1;如果是服务端批量推理,导 batch=4 或 8,配合动态 batch 一起用。batch 越大,单张图的理论吞吐越高,但首次推理的延迟和显存占用也越高。这里有个经验值:消费级显卡上,YOLOv5s 在 640 分辨率下,batch=1 时显存占用大概 1GB 出头,batch=8 时接近 5GB,这不是线性关系,因为中间层激活值的峰值占用会随 batch 放大。所以别一上来就导 batch=32,先按实际需求来。

3.3 opset 版本:高了低了都出事

opset 是 onnx 的算子集版本,它决定了计算图里能用哪些算子、这些算子的语义是什么。默认值一般是 12,这个版本是个比较稳的甜点。往低了走,某些算子不被支持,导出时 PyTorch 会报“不支持的算子”然后中断;往高了走,导出的图里可能出现较新的算子,你下游的推理引擎不一定认识,表现为加载时报算子不支持。

所以 opset 的选择原则是就低不就高,以下游运行时的支持上限为准。如果你确定下游用的是某个具体版本的推理引擎,去查它支持的 opset 范围,取那个范围内的一个偏低值。如果没有明确的下游约束,用默认值 12 就行,不要因为“版本新看起来更先进”去调高。

还有一个不太常见的坑:某些 PyTorch 版本在导出时会因为内部注册的符号函数变化,导致同一个 opset 下导出的图结构不同。这就是为什么前面强调要把 torch 版本也钉死——同样的权重、同样的命令、不同的 torch,导出的 onnx 可能不完全一样。

3.4 --simplify 与 --dynamic:什么时候开,什么时候别碰

--simplify会调用 onnx-simplifier 做常量折叠、冗余节点消除和形状推断,通常能让图更干净、推理更快。建议默认开启,但要知道它的两个副作用:一是会增加几十秒的导出时间,二是它在某些动态形状的图上会简化失败。失败时的表现不太友好,有时是抛异常,有时是静默跳过——你以为简化了其实没有。所以我一般在导出完成后会对比一下开启和关闭简化时生成的文件大小,差异明显才算真的生效。

--dynamic打开后导出的 onnx 会把 batch 维和空间维标记为动态,好处是一份文件可以吃不同尺寸的输入,坏处是推理引擎无法针对固定尺寸做极致优化,性能通常比固定形状低一截。更要命的是,YOLOv5 的检测头里有基于特征图尺寸动态生成 anchor 网格的逻辑,动态尺寸下这部分必须在运行时重新计算,如果下游的推理框架不支持这种动态算子,整个图就跑不起来。

我的判断标准很直白:部署端输入尺寸固定,就别开--dynamic,老老实实导固定形状,性能最好;只有在“一份模型要服务多种分辨率”这种真实需求下才开,并且开之前一定要和下游确认他们的运行时支持动态形状。

3.5 导出之后的三重校验,缺一不可

文件生成了不等于正确。我固定做三步校验:

第一步,结构校验。用 onnx 官方库把文件加载一遍并做模型检查,这一步能挡住大部分结构性错误,比如图里有悬空节点、输入输出类型不合法、权重张量维度对不上。

import onnx m = onnx.load("runs/export/onnx_fp32_b1/best.onnx") onnx.checker.check_model(m) print("inputs:", [(i.name, [d.dim_value for d in i.type.tensor_type.shape.dim]) for i in m.graph.input]) print("outputs:", [(o.name, [d.dim_value for d in o.type.tensor_type.shape.dim]) for o in m.graph.output]) print("opset:", m.opset_import)

打印输入输出的形状这一步别省,很多人就是因为没看,结果部署端按错误的形状喂数据,一路错到底。输出形状应该是[batch, 候选框数, 5+类别数],如果类别数是 80,那就是 85 列。

第二步,运行时校验。用 onnxruntime 加载并跑一次随机输入,确认能正常前向、输出没有 NaN。这一步能发现算子不被支持、精度溢出这类问题。

import numpy as np, onnxruntime as ort sess = ort.InferenceSession("runs/export/onnx_fp32_b1/best.onnx", providers=["CPUExecutionProvider"]) name = sess.get_inputs()[0].name x = np.random.rand(1, 3, 640, 640).astype(np.float32) y = sess.run(None, {name: x}) print("output shape:", y[0].shape) print("has nan:", np.isnan(y[0]).any(), "range:", y[0].min(), y[0].max())

第三步,端到端校验。直接用仓库自带的推理脚本加载这份 onnx 去跑几张真实图片,看框的位置和类别是否正常。这一步才是真正的验收标准,因为前两步只能证明“文件是合法的”,不能证明“结果是对的”。

python detect.py --weights runs/export/onnx_fp32_b1/best.onnx \ --source samples/ --imgsz 640 640 --conf-thres 0.25 --device 0

三步都过,这份 onnx 才算能交付。

4. 导出 TensorRT engine:真正吃硬件的那一步

4.1 export.py 里 engine 那段到底做了什么

engine 导出走的是另一条完全不同的路径。脚本会把刚才的模型转成 onnx 中间态(在内存里或临时文件),然后用 TensorRT 的构建器加载这份图,配置构建参数,最后执行一次耗时的编译过程,产出一个序列化的执行计划文件。

这个编译过程里,TensorRT 会做几件事:把连续的卷积、批归一化、激活函数融合成一个复合层,减少内核启动开销;在开启半精度时逐层判断能否用 FP16 计算而不损失太多精度;针对当前显卡的 SM 数量、寄存器、共享内存容量,从内置的候选内核里自动挑最快的那一个进行实测对比。这几件事加起来就是为什么第一次构建这么慢——尤其是内核自动调优,它会为每一个层试跑多个候选实现。

理解这一点就能明白:engine 构建是一次线下编译,不是一次转换。所以它慢是正常的,慢不代表卡死。

4.2 FP32、FP16、INT8 三档精度的收益账

--half打开就是 FP16。这是绝大多数场景的推荐选择:显存占用大约减半,推理速度在主流消费级显卡上通常有 1.5 到 2 倍的提升,精度损失在检测任务上一般体现在小数点后一位的 mAP 上,对于 0.5 以上的置信度阈值来说基本感受不到。

FP32 是不开--half的默认状态,精度最稳,适合做基线对比,但速度和显存都不划算,实际部署很少用。

INT8 需要额外的东西:一份校准数据集,通过在真实数据上跑一遍前向、统计每一层激活值的分布,生成一张缩放系数表,构建时用这张表把 FP32 的权重和激活量化到 8 位整数。收益是显存再砍一半、速度再提一截,代价是精度损失比 FP16 明显,而且校准数据的分布必须和实际场景接近,否则量化误差会放大。导出脚本本身不负责校准,校准流程是和验证脚本配合完成的,具体调用方式在不同版本里差别不小,务必以当前版本的参数说明为准,别照抄网上的老命令。

我的建议是:先导 FP16,用它跑一遍完整验证集看 mAP;如果速度还不够,再考虑 INT8,并且一定用真实场景数据做校准和回归。

4.3 workspace 给多少:小了构建失败,大了吃显存

--workspace参数控制 TensorRT 构建阶段可以使用的临时显存上限,单位是 GB。这个值不是给推理用的,是给编译期做内核调优用的——候选内核需要一个草稿纸来试跑,草稿纸太小,某些需要大临时空间的优化就做不了,构建器会退而求其次选一个慢一点的实现,甚至在极端情况下直接构建失败。

默认值一般是 4GB,对于 YOLOv5 这个规模的模型,大部分情况下够用。但如果你的输入分辨率很大(比如 1280 以上)、batch 也大,4GB 可能不够,表现是构建时间异常长或者报显存相关的错误。这时候可以往上调到 8GB 甚至更大。但要注意:这个峰值是构建期占用的,构建完成后就会释放,所以不要因为“显卡只有 8GB 显存”就不敢设 8。真正需要注意的是构建期间别让别的进程占着显存。

反过来说,如果你的显卡本来就只有 4GB 或 6GB 显存,把 workspace 设得比显存还大也没意义,构建器申请不到。这种情况下正确做法是降低 batch 或者降低分辨率,而不是硬撑 workspace。

4.4 动态 shape 的 profile 设置与踩坑点

engine 导出时如果开了--dynamic,构建器必须知道输入张量的最小、最优、最大三个形状,因为不同形状下最优的内核可能不一样,TensorRT 会为这三档分别做优化。脚本里通常会自动生成一组 profile,比如最小 batch 为 1、最优为导出时指定的 batch、最大为它的两倍左右。

这里有两个高频问题。第一个是最大形状设小了,推理时喂了超范围的 batch 或分辨率,直接抛“输入尺寸超出 profile 范围”的异常。解决办法是把最大值设得比你的实际需求留出余量。第二个是动态空间维(宽高)在某些下游封装里不被支持,因为很多部署框架假定输入形状固定,遇到动态维会在内存分配上出错。如果你用的是自己写的前后处理链路,动态形状问题不大;如果是接现成的推理服务框架,先确认它是否支持动态形状,再决定要不要开。

我个人的经验是,工程上 90% 的场景输入尺寸都是固定的,动态更多是给“同一份模型跑多种分辨率做模型集成”这类需求准备的。别为了“灵活”两个字给自己加调试成本。

4.5 构建耗时:怎么区分“慢”和“挂”

一次 YOLOv5s 的 FP16 engine 构建,在中端显卡上大概几分钟,在高端卡上可能一两分钟,在低端卡或开了动态形状、大 batch 的情况下十几分钟也正常。终端上一般会打印进度信息,比如正在优化第几层。

判断是不是挂掉,看两个信号:一是显存占用有没有在动,如果显存一直稳在一个值、CPU 也几乎不占,那大概率是真卡住了;二是看终端最后一行停在哪个阶段,如果停在某个层名反复不往前走,多半是那一层触发了某个编译器的边界情况。卡住时最有效的排查手段是换一个更小的分辨率、更小的 batch 再试一次——如果小规模能过,说明是资源问题;如果小规模也卡在同一层,说明是那一层的算子有问题,可以考虑升级或回退 TensorRT 版本。

5. 实测故障排查:几个我真实遇到过的报错链路

5.1 onnx 导出成功但推理结果一团乱

这个现象最典型,表现是框的位置整体偏移、大小离谱,或者类别置信度普遍偏低。排查顺序我固定按下面走。

第一看输入预处理是否一致。训练和导出用的是 letterbox 填充,保持长宽比、用灰色填充到目标尺寸。如果部署端用了简单的拉伸缩放,输入分布就变了,输出必然偏。这一点在宽高比差异大的图上(比如长条形的监控画面)特别明显。

第二看颜色通道顺序和归一化。仓库默认读图是 BGR 顺序,归一化是除以 255。如果部署端用 RGB 直接喂进去,红蓝通道互换,检测结果会明显变差但不会全错,这种“差一点点”的情况最难发现。归一化漏了更明显,输出会完全失去意义。

第三看输出解析。原始输出是[1, 25200, 85],前 4 个是中心点坐标和宽高(相对于网格的值),第 5 个是目标置信度,后面是类别分数。解析顺序错了,框就会全部集中在图像左上角或者尺寸趋近于零。

一个快速定位的技巧:拿一张图,同时用.pt和.onnx通过同一个推理脚本跑一遍,把两边的原始输出张量存下来做逐元素比对。如果数值差异在 1e-3 量级,说明导出没问题,问题在预处理;如果差异巨大,说明导出这一步本身有问题,回头查 opset 和简化选项。

5.2 构建 engine 时报某个层不支持

这类报错信息通常比较具体,会告诉你哪个算子、哪个节点不被支持。根因往往不是“模型有问题”,而是算子集版本和 TensorRT 版本的匹配问题。高版本的 onnx 可能引入了较新的算子,而当前 TensorRT 版本还没实现它。

处理方式有两个方向。一是把--opset降下来重新导 onnx,让计算图用更基础的算子表达;二是升级 TensorRT 到支持该算子的版本。前者风险小、见效快,后者要注意整条工具链的兼容性。我一般优先降 opset,因为它的影响面只在这份图上,不会波及现有环境。

还有一种情况是自定义算子。如果你在训练时改了检测头或者加了特殊模块,导出的图里就可能出现 TensorRT 不认识的自定义算子。这种情况没有捷径,要么用插件机制把它实现成 TensorRT 的自定义层,要么在导出前把这段逻辑挪到后处理里去。

5.3 PyCUDA 缺失导致的连锁失败

症状是 onnx 导出成功、engine 导出在编译阶段抛出一个和显存指针相关的异常。很多人会顺着异常信息去查显存不足,其实是 PyCUDA 没装上或者版本不匹配。

判断方法很简单,跑一下前面自检脚本里的第三条,如果import pycuda.autoinit就报错,那根因就找到了。修复方式是装一个和当前 CUDA 版本匹配的 PyCUDA,注意这个包对 CUDA 版本很敏感,装错了会在导入时报找不到动态库。

顺带提一个隐蔽的变体:PyCUDA 装了,但和 PyTorch 使用的 CUDA 上下文冲突。表现是 engine 构建到一半报“无效的上下文”。这种情况的解法是在导出脚本里确保 PyTorch 的 CUDA 初始化和 PyCUDA 的上下文不要交叉使用——通常是通过延迟导入或者把两者放在不同的进程里。

5.4 engine 换机器加载失败

这是最不该踩但踩的人最多的坑。前面说过 engine 和硬件、软件版本绑定,所以换机器前先问三个问题:目标机器的显卡架构和构建机器是否一致?TensorRT 大版本是否一致?驱动版本是否满足要求?

只要有一个不满足,加载时就会报反序列化失败。报错信息通常很笼统,只说“无法反序列化”或者“引擎版本不兼容”,不会告诉你具体差在哪。所以正确的交付方式是:给引擎文件的同时,附上构建环境的完整描述,包括 TensorRT 版本、CUDA 版本、构建时的显卡型号、输入形状配置。对方拿到这些信息才能判断能不能直接复用。

如果一个项目需要部署到多种硬件,那就不要交付 engine,交付 onnx 加一份构建脚本,让每个目标的部署环境自己编译一份。虽然多花了编译时间,但省掉了所有兼容性扯皮。

6. 把导出做成可重复的工程动作

6.1 写一个批量导出脚本,别手敲命令

权重一多,手敲命令行就变成了体力活,而且容易敲错。我通常写一个简单的 shell 脚本,把权重列表、分辨率、batch 组合遍历一遍。

#!/usr/bin/env bash set -euo pipefail WEIGHTS_DIR="runs/train" OUT_DIR="runs/export" IMGSZ_LIST=(640 640) BATCH_LIST=(1 4) for w in "${WEIGHTS_DIR}"/*/weights/best.pt; do tag=$(basename "$(dirname "$(dirname "$w")")") for b in "${BATCH_LIST[@]}"; do echo ">>> export ${tag} batch=${b}" python export.py \ --weights "$w" \ --include onnx engine \ --imgsz "${IMGSZ_LIST[@]}" \ --batch "$b" \ --opset 12 \ --simplify \ --half \ --device 0 \ --project "${OUT_DIR}" \ --name "${tag}_b${b}" done done

set -euo pipefail这三件套别省,它能让脚本在任何一个命令失败时立刻停下,而不是带着错误继续跑下去,最后产出一堆半成品文件却不知道哪一步出的问题。批量导出最怕的就是这个——看着跑完了,实际上一半是坏的。

脚本跑完之后我会再补一个校验循环,对每个生成的 onnx 做一次结构检查和一次随机前向,把不通过的挑出来。

import glob, onnx, numpy as np, onnxruntime as ort for f in sorted(glob.glob("runs/export/**/*.onnx", recursive=True)): try: m = onnx.load(f) onnx.checker.check_model(m) s = ort.InferenceSession(f, providers=["CPUExecutionProvider"]) shape = [d.dim_value or 1 for d in s.get_inputs()[0].type.tensor_type.shape.dim] y = s.run(None, {s.get_inputs()[0].name: np.random.rand(*shape).astype(np.float32)}) print("OK ", f, y[0].shape) except Exception as e: print("FAIL", f, repr(e))

6.2 产物命名和版本记录,比导出本身更重要

文件名里必须能一眼看出关键信息。我的命名规则是模型名_分辨率_batch_精度_日期,比如det_v3_640_b1_fp16_20240612.engine。这样半年后回来看到文件,不用翻日志就知道它是什么配置的产物。

除了文件名,每次导出往一份记录表里追加一行:权重路径、权重文件的哈希值、导出时间、环境版本三元组(torch / TensorRT / CUDA)、导出参数、产物大小、验证结果。这张表看起来麻烦,但它是出问题时唯一能救命的东西——尤其当你要回答“为什么上个月导出的那份精度比现在这份高”这种问题时,没有记录就只能重新跑一遍全部实验。

哈希值这一步不要省。权重文件被覆盖、被误传到别的分支、被重新训练写回同一个路径,这些在多人协作里都是常事。有了哈希,你才能确认“这份产物确实是从这份权重来的”。

6.3 用推理脚本做端到端回归,而不是只看文件大小

最后一步是回归。我的做法是拿一份固定的验证图片集(十几张就够,覆盖不同场景、不同光照、不同目标密度),对每一份新产物跑一遍推理,把检测结果数量和置信度分布记录下来,和历史值对比。如果数量突变或者置信度整体偏移,立刻回头查。

这一步能挡住很多“文件合法但结果不对”的问题。比如简化选项导致某个算子被错误折叠、动态形状配置导致网格计算偏差、半精度在某些层上溢出,这些都通不过端到端回归,但都能通过前面的结构校验。

在工程实践里,我还会准备一份.pt的推理结果作为基准,把 onnx 和 engine 的输出框拿去和它做 IoU 匹配,统计匹配率和平均 IoU。FP16 的 engine 通常能达到 99% 以上的匹配率,如果掉到 90% 以下,就说明精度损失不可接受,得回头调整精度配置。

我个人在多次交付里最深的体会是:导出这一步的技术难度其实不高,真正费时间的是“确认导出结果是对的”这件事。前面那些参数、版本、校验步骤,本质上都是在为这个确认过程服务。所以与其在报错时到处搜,不如一开始就把校验链路搭起来——把自检脚本、结构校验、随机前向、端到端回归这四步固化成一个小工具,每次导出后跑一遍,十几秒钟就能给你一个明确的是或否。这个投入在第二个项目上就回本了。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进