项目排期表上Phase A · Step 3模型导出 ONNX只占了一个不起眼的小格子。当时团队预估这是一个顶多两三天的技术活结果我前后折腾了将近一周才敢说这个导出是稳的。后来复盘发现问题不在于torch.onnx.export这个接口本身有多难而是整个导出链路里的隐藏依赖、版本矩阵和下游约束远比大多数人想象的多。这篇内容我围绕一个具体场景来写某视觉检测模型训练基于 PyTorch目标产物是能在服务端CPU环境稳定运行的ONNX文件后续推理引擎选用ONNX Runtime。如果你也正处在模型训练完、准备交付到部署环境这个阶段那这篇应该能帮你少踩不少坑。1. 为什么导出ONNX值得单独拆成一个Step1.1 项目里这条部署链路ONNX到底卡在哪一环很多同学第一次接触ONNX是在某个快速上手教程里跟着代码三五行就把模型转出来了一个 .onnx 文件。于是大家天然觉得这玩意儿不是有手就行吗项目排期里为什么要给它单独开一个Step实际上在真实项目里ONNX导出是训练和部署之间的交接仪式。训练侧模型是活的有动态Batch、有训练状态、有各种中间变量它还可以继续训练、继续调参。部署侧模型是死的输入尺寸固定、计算图固定、运行环境完全独立。ONNX就是这个由活到死的转换点。在这条链路里ONNX承担的是中间表示IR的角色。它本身不负责高效推理它的价值是让上层训练框架PyTorch等和下层推理运行时或硬件编译器解耦。一个 .onnx 文件导出之后训练框架里的各种花活就彻底与它无关了下游只认这个静态的描述文件。这样换来两个核心能力一是同一份模型可以在不同运行时上运行二是在不依赖深度学习框架环境的前提下完成性能优化。1.2 不导出的代价原框架推理的三个现实约束我自己经历过先跑通再说的阶段当时方案是直接在 PyTorch 环境里加载模型用原框架做推理服务。这个方案看起来省事但真正推进到项目后期会撞上三面墙。第一环境锁死。服务器上要装对应版本的 PyTorch、CUDA 库、一堆传递依赖升级一个系统库都可能让模型服务瘫痪。模型交付出去基本等于把整套训练环境也交付出去运维成本直接起飞。第二推理效率上不去。训练框架的模型执行方式更关注灵活性同样的算子框架内建执行逻辑往往比专门的推理引擎要多一层调度开销。尤其在高并发、小Batch场景下这个差距会非常明显。第三硬件适配能力弱。如果后续业务要上推理加速卡或者边缘设备这些硬件厂商的编译器通常只吃 ONNX 或自家格式不支持直接加载 PyTorch 的模型权重。所以导出 ONNX 不是可选项而是把模型从训练链路里解放出来、让它进入工程化体系的必经环节。1.3 ONNX导出前后的责任边界导出前你的身份是算法工程师关注的是模型精度、Loss 曲线、效果指标。导出后你的身份要切换成部署交付者关注的是算子兼容性、图优化空间、推理延迟。很多人都没意识到这个身份切换意味着你不能再抱着一颗算法模型完美就行的心态。从项目管理的角度看单独把模型导出拆成一个Step本质上是强制做一次交付物确认。在排期里明确它能逼着团队把一些原本含糊的问题提前暴露模型结构是否确定算子是否会有兼容性风险部署环境的 ONNX Runtime 版本是什么这些问题如果留在部署阶段再发现排查成本会呈指数上升。2. 动手导出之前模型冻结、算清单与版本对齐2.1 模型处于eval且冻结状态这是最基础也最容易被忽略的一件事。如果模型里有 BatchNorm 或 Dropout导出前必须确认已经切换到 eval() 模式。否则BatchNorm 层在训练模式下会用到当前 batch 的统计数据而不是训练好的 running_mean 和 running_var导出后的模型推理结果可能会漂移。实际操作里我的习惯是不仅调用 model.eval()还会顺带看一眼模型里各层的 running_mean 有没有更新。尤其是从检查点恢复的模型如果中间出现过训练中断、继续训练、再保存这类操作一旦本地变量和 state_dict 对不上eval() 也不一定能兜住。另外冻结不等于eval。冻结通常指所有参数 requires_gradFalse这样既能避免误改权重也能让后续 trace 出的计算图更干净。我见过一些项目导出时忘了关 requires_grad导是能导出来但文件里会额外带上梯度相关的节点白白增加图体积。2.2 输入输出的shape和dtype提前定死ONNX文件里最核心的元信息就是输入输出的 shape、dtype。虽然 ONNX 支持动态维度但动态维度并不是随便填那么好用下游推理引擎对动态shape的实现差异很大。所以导出前需要确定一件事这个模型上线后输入尺寸是固定的还是需要支持任意大小如果业务上输入尺寸固定比如检测模型统一 Resize 到 640x640 再进网络那就直接在导出时把输入样例定成 (1, 3, 640, 640)让 ONNX 导出逻辑把中间所有 shape 全部算死。这样做出来的模型在推理引擎里图优化空间最大速度也最稳。如果确实需要动态比如第一维 batch 可变或者 NLP 类的序列长度可变那就必须明确声明动态轴。这里有个经验ONNX 的动态轴支持是选择性动态不是全局动态。如果你把某个维度标成动态但图内部有一些算子偷偷假定这个维度是固定值导出时会报错或者导出的模型行为不正确。因此动态轴不要贪多能用固定 shape 就用固定 shape。2.3 opset版本与推理引擎的对应关系opsetOperator Set是 ONNX 的算子版本集合。每次 ONNX 规范更新都会增加新的算子、修正旧算子定义。PyTorch 导出时的 opset 参数决定它用哪个版本的算子定义去序列化计算图。这个版本号有两端必须对齐一端是 PyTorch 导出侧能否正常生成另一端是 ONNX Runtime 或硬件编译器能解析到什么版本。我遇到过最典型的问题本机 PyTorch 版本较新随手填了 opset17结果目标环境的 ONNX Runtime 只支持到 opset 15模型加载直接报错。一个稳妥的做法是在导出前先确认目标推理引擎支持的 opset 上限然后导出时的 opset 既不要低于 11太老会有很多算子缺失也不要高于目标引擎的上限。比如目标 ONNX Runtime 版本支持 opset 15那导出就填 15别为了新而去填更高的版本。3. trace、script还是dynamo三种导出路径怎么选3.1 torch.onnx.export默认走tracePyTorch 里最经典的导出接口是 torch.onnx.export大多数人知道它内部默认用 tracing 机制给定一份示例输入把模型前向过程中实际执行到的算子记录下来形成一个静态图。trace 的好处是简单直观只要前向逻辑是直来直去的几乎一条命令就能搞定。但它的局限性也很明显只能导出这次执行走过的路径。如果你的模型里有一段 if 分支而示例输入恰好只走了其中一条分支那另一条分支在导出的模型里压根不存在。后续推理遇到另一条分支的输入结果就是错的。我在某个分割模型里就踩过这个坑。模型里有一个针对不同分辨率的判断分支导出时我用的示例输入是 640 分辨率成功导出了。结果真跑到 1280 分辨率数据上模型直接乱输出。排查半天才发现trace 过程把分辨率判断里的较小分支裁剪掉了。3.2 script方式适合控制流为了解决 trace 丢失控制流的问题可以用 torch.jit.script 先把模型转成 ScriptModule再导出 ONNX。script 走的是解析源码的方式显式保留 if、for 等控制流结构图里会生成相应的控制流节点。但 script 对模型代码有要求不能随意使用 Python 的动态特性。你写模型时可能觉得用 .item() 取个数值很方便列表推导式很 Pythonic这些到了 script 阶段全都可能报错。所以如果你的模型并不复杂也没有明显的动态分支我一般不建议为了用 script 而用 script改写代码的时间成本可能超过收益。3.3 dynamo新路径的实际体验PyTorch 2.x 之后推出了 dynamo 相关能力torch.onnx.export 也提供了 dynamoTrue 的导出选项。我实际试用过几次最大的感受是它对 Python 代码的兼容性更好很多原来需要手工改写的代码模式可以原样通过。图结构和 opset 支持也更规范。但目前来说dynamo 导出仍然有一些算子覆盖上的坑尤其是一些自定义的 C extension 算子或者某些用了高阶梯度能力的模型。我的建议是新项目可以优先尝试 dynamo但如果测试中出现解释不清的行为差异就退回传统 trace 再看看。4. 完整导出脚本与产物级验证4.1 最小导出脚本下面这个脚本是我在项目里反复用过的基础模板参数都标注了用途import torch model load_model_from_checkpoint(best.pt) model.eval() # 冻结参数避免误操作带梯度的节点 for p in model.parameters(): p.requires_grad False # 输入样例shape (batch, channel, height, width) dummy_input torch.randn(1, 3, 640, 640) torch.onnx.export( model, dummy_input, model.onnx, opset_version15, input_names[input], output_names[output], dynamic_axes{ input: {0: batch_size}, output: {0: batch_size}, }, do_constant_foldingTrue, )有几个细节值得展开。dynamic_axes 里我通常只开放 batch 维。如果你的业务场景里宽高也变化可以加上 {2: height, 3: width}但要注意 resizing 类操作对动态尺寸的算子兼容性。do_constant_foldingTrue 默认开启它会把一些只依赖常量参数的子图预先计算出来比如 Conv 层的 bias 加法。开启后模型文件可能更小推理时也能少跑几个节点建议保持开启。4.2 onnx.checker和shape_inference导出完成后别急着拿去上服务。第一步先用官方工具做基础体检import onnx model onnx.load(model.onnx) onnx.checker.check_model(model) # 重建图内的shape信息补全中间张量的维度 model onnx.shape_inference.infer_shapes(model) onnx.save(model, model.onnx)onnx.checker.check_model 检查的是模型的格式合法性包括节点输入输出是否匹配、属性是否合法、图结构是否完整。这属于语法检查不保证语义正确但能拦下大量低级问题。shape_inference 是我每次导出后必跑的一步。很多框架导出的图里中间张量的 shape 信息是缺失的推理引擎拿到这种图往往只能按动态思路处理性能受影响。手动跑一次 shape inference相当于给图补全了尺寸标注下游引擎的优化空间会大很多。4.3 用ONNX Runtime做数值一致性验证格式检查通过只是第一步真正的考验是数值一致性。我的做法是把原始 PyTorch 模型的输出和 ONNX 模型的输出做逐元素对比import numpy as np import onnxruntime as ort import torch # 固定输入多次采样求最大误差 x torch.randn(8, 3, 640, 640).numpy() sess ort.InferenceSession( model.onnx, providers[CPUExecutionProvider], ) outputs_onnx sess.run(None, {input: x}) with torch.no_grad(): model.eval() outputs_torch model(torch.from_numpy(x)).numpy() # 统计最大绝对误差和相对误差 max_diff np.max(np.abs(outputs_onnx[0] - outputs_torch)) print(max abs diff:, max_diff)这里要注意PyTorch 模型和 ONNX Runtime 模型在浮点运算顺序上可能有小差异比如 Convolution 的 im2col 和矩阵乘的累加顺序不同。所以不要期望二者 diff 为 0一般观察 1e-4 到 1e-5 量级的最大误差即可。如果误差到了 1e-2 甚至更大那就要怀疑导出过程中某个算子被错误降级或者动态分支被裁掉需要往具体层去定位。5. 这步最容易翻车的五个实操坑5.1 动态轴配了但输入样例写错这个坑极其隐蔽。某个模型的导出脚本里声明了 height 和 width 是动态轴但 dummy_input 用的是 (1, 3, 224, 224)。文件导出来一切正常也能跑但如果你用 (1, 3, 512, 448) 去测试中间某个 Resize 节点可能计算出不符合预期的 shape甚至整个推理崩溃。原因在于部分算子对输入尺寸有内部假设trace 时这个假设被写死在了常量节点里。最典型的例子就是某些 F.interpolate 的 scale_factor 被转成了固定数值一旦输入尺寸变化插值目标尺寸就错了。对付这个问题的排查方法是导出后用多组不同 shape 的输入跑一遍 ONNX Runtime看输出 shape 是否符合预期而不只测一组。5.2 upsample算子被rewrite语义分割类模型经常会在解码器里用到上采样。PyTorch 里的某些上采样实现在导出 ONNX 时会被转成一系列自定义节点例如 Resize 配合 coordinate_transformation_mode 的复杂组合。不同 opset 版本下这个转换结果差别很大。我在一个分割模型上遇到过PyTorch 里用 F.interpolate(modebilinear)导出为 opset 11 时一切正常换成 opset 15 后ONNX Runtime 提示 Resize 节点的某个属性不支持。最后没办法只能强行把导出 opset 降回 11或者改模型内部的上采样方式比如换成 PixelShuffle 或卷积堆叠。这提醒我们opset 不是越新越好要结合算子实际转换结果来定。如果某一次切换 opset 版本后出现奇怪报错先把版本回退对比一下模型输出很多时候问题就消失了。5.3 数值对不上数值对不上的原因很多最常见的是预处理不对齐。ONNX Runtime 输入是裸 Tensor它不关心你在 PyTorch 里是怎么做归一化的。如果你在模型内部做了 Normalize 然后用均值和方差做了减法那 OK但如果 Normalize 在模型外部用 PIL 或 OpenCV 做的部署时也得复刻同样的顺序、同样的通道格式、同样的缩放。我遇到过团队里前后端联调大半天最后发现是 HWC 和 CHW 的通道顺序问题。PyTorch 模型内部是 CHW预处理端输出的是 HWC 且忘记 transpose结果推理结果的准确率断崖下跌。5.4 图优化选项导致显存爆炸ONNX Runtime 里有一个 graph optimization level 配置常见的有 ORT_DISABLE_ALL、ORT_ENABLE_BASIC、ORT_ENABLE_EXTENDED、ORT_ENABLE_ALL。这个配置能自动融合算子、重排节点、消除冗余理论上推理性能会更好。但你调高优化级别后要重新做一轮精度验证。我遇到过把优化级别调到 EXTENDED 之后某层被融合成了新的算子组合浮点误差变大最终线上效果出现可见的劣化。还有一个极端情况是开启闭环优化后某些大模型推理时内存占用反而上升。所以我的建议是性能调优时优化级别从 BASIC 开始往上试每提升一档都回归一次精度和延迟。不要图省事直接拉满 ALL。5.5 opset太高导致推理引擎不支持这算是版本问题里最高发的。通常是开发环境用的 PyTorch 很新导出默认 opset 也高代码一写就导出成功。等到部署环境加载时报类似Model opset version X is not supported的错误。这种问题的排查链路很简单先看报错信息里要求的 opset 上限再看导出脚本里的 opset_version改了之后重新导出。关键是不要在本地换 PyTorch 版本因为升级或降级 PyTorch 可能影响模型权重加载行为正确做法是固定 PyTorch改 opset。6. 导出完成之后接入运行时和后续迭代流程6.1 session.options里值得调整的参数ONNX Runtime 跑起来之后有几个参数值得认真调一下。首先是 intra_op_num_threads 和 inter_op_num_threads分别控制单个算子内部的线程数和不同算子间的并行线程数。CPU 推理场景下这两个值不是越大越好设得比物理核数还多是常见性能杀手。其次是 execution_mode默认是 ORT_SEQUENTIAL如果模型里没有明显的并行分支我建议保持默认。并行模式在有分支结构的模型上有收益但也会引入额外的线程调度开销小模型反而变慢。另一个容易被忽略的是 session 的 graph_optimization_level。我上面提过优化级别会影响精度正式部署前应该把这一项固定下来并写进配置文档避免不同环境加载同一份模型时行为不一致。6.2 把导出-验证-提交固化到pipeline里模型迭代不是一次性的项目到了中后期可能一天要导出好几次新版本模型。如果每次都靠手动执行脚本很容易在不同的环境里导出结果不一致。我在项目里做了一件事把导出脚本、shape inference、数值一致性验证、样例输出检查打包成一个自动化脚本。每次训练出新的权重文件只要执行一条命令就能得到一个确认可交付的 ONNX 文件加上一份校验报告。这个流程里印象最深刻的一点是校验环节里除了数值比对还固定了几组真实业务输入用来确认模型在常见输入下的输出是否合理。数值比对只能保证和 PyTorch 一致但不保证符合业务预期尤其当 PyTorch 模型本身就有点问题的时候容易把错误一套搬到 ONNX。所以我建议在自动化流程里加一个业务回归检查步骤比如固定几张图输出指标高于某个阈值才算通过。6.3 下次迭代时的注意点模型结构一旦有变化ONNX 导出不是简单重跑一遍就行。要特别关注新增的模块是否引入了不支持导出的操作。我的经验是在模型编码阶段就提前审查每个新算子的可导出性比导出时再发现要省力得多。另外模型文件在多次导出的过程中可能会积累一些旧版本残留图节点。如果导出的模型体积比预期大出不少可以用 onnxsim 这类工具做一次简化。它能把常量折叠掉、移除冗余节点把图整理得更清爽。但用完 onnxsim 之后也要重新做一遍精度验证确保简化过程没有引入语义变化。回头看这个 Step我的直观感受是真正难的不是那一行导出代码而是围绕导出建立起一套可重复、可验证、可追溯的流程。当模型从训练环境迈向生产环境的那一刻它能跑是不够的你得确保它能稳定地、和预期一致地跑很久。这一步做扎实了后面接入推理服务、性能调优都会顺很多。