1. 导出前的准备工作与方案选型做模型导出 ONNX 这事乍一看好像就是把torch.onnx.export跑一遍的事但实际上十个模型九个坑前面准备不到位后面导出、验证、上板子全都是重新返工的节奏。我这次做的是 Phase A 第 3 步目标是把训练好的 PyTorch 模型导成 ONNX为后续的量化、推理加速和端侧部署铺路。整个流水线里ONNX 导出算是最关键的中间环节——它卡在“训练”和“部署”之间导不好后面所有步骤都是空中楼阁。先说准备工作。这个阶段看起来是“准备”实际上直接决定了后面能不能少折腾。第一件事是锁版本。PyTorch 版本、ONNX 版本、onnxruntime 版本这三者必须匹配。我见过不少网友在群里截图报错Failed to export an ONNX attribute onnx::Gather这类问题翻来覆去看了半天最后发现是 torch 1.8 和最新版 ONNX 不兼容。我的建议是PyTorch 用 1.13 以上最好是 2.xONNX 库用 1.13 以上onnxruntime 用最新稳定版。这三个版本组合基本能覆盖绝大多数模型的导出需求。如果你用的是老框架遇到算符缺失的问题第一反应不是改代码而是看看版本是不是太旧了。第二件事是确定导出协议。ONNX 导出有两套逻辑一套是 trace追踪一套是 script脚本化。PyTorch 默认走的是 trace 机制也就是给定一组示例输入把计算图里的实际执行路径记录下来。这意味着如果你的模型里有if data.sum() 0:这类依赖数据值的分支trace 只会记录满足当前输入条件的那一条路径另一个分支永远不会出现在 ONNX 图里。我之前导出过一个带 attention mask 的模型就是因为 mask 恰好全为 1trace 阶段把 mask 分支直接“优化”掉了导出来的模型在真实场景下结果完全错误。所以使用 trace 之前要么保证模型结构是静态的要么用torch.jit.script把动态分支显式 script 化要么干脆在设计模型时避免数据依赖的分支。第三件事是输入输出的命名规划。这个容易被忽略但到了部署阶段不管是 onnxruntime 还是 TensorRT都要按名字取输入输出。如果导出时没有显式指定input_names和output_namesONNX 会自动生成input.1、output.1这种名字你后面的推理代码就会变成魔法字符串满天飞。我习惯在导出前就把输入命名为input、输出命名为output如果模型有多个输入就在名字里加后缀比如input_ids、attention_mask一目了然。1.1 动态轴与静态轴怎么选动态轴dynamic axes和静态轴static axes的选择是我在导出前想得最多的事。所谓动态轴就是允许模型在推理时接收不同尺寸的输入。比如一个目标检测模型训练时输入图片是 640x640但部署时用户上传的图片五花八门这时候如果不做动态轴ONNX 模型就固定接收 640x640 的输入其他尺寸一律报错或强制 resize灵活性大打折扣。但动态轴不是免费午餐。动态轴会在模型图里引入额外的 Shape 计算和 Gather 算符这会带来不小的性能开销尤其是在 TensorRT 这类对静态 shape 做了深度优化的推理引擎上动态 shape 往往导致 TensorRT 回退到 fallback 模式推理速度直接掉一大截。所以我的建议是能静态就静态必须动态再动态。如果是服务于固定输入尺寸的推理服务比如视频流分析、固定分辨率检测静态轴是最稳妥的选择如果是给开放平台用的通用模型那就老老实实配dynamic_axes。举个例子我这次导出的模型输入是[1, 3, 640, 640]我用了动态轴只作用于 batch 维度和高度宽度维度代码长这样dynamic_axes { input: {0: batch, 2: height, 3: width}, output: {0: batch, 2: height, 3: width} }注意这里我没有把通道维度设为动态。因为模型结构里的卷积层权重是[64, 3, 3, 3]这种固定 shape通道数变化会引起卷积权重不匹配动态通道没有任何实际意义。类似的教训是设定动态轴时只设定确实会变化的维度不要想着“全都动态”那样会引入大量无用的 Shape 操作。1.2 模型锁定的重要性还有一个极其容易踩的坑导出前模型必须处于 eval 模式并且锁住参数。很多人在训练完直接调用 export忘记model.eval(),结果导出来的 ONNX 图里带着 dropout 层、batchnorm 层的训练行为推理时表现为输出随机抖动、精度不稳定。而且如果在导出前没有torch.no_grad()包裹模型的前向过程还会计算梯度不仅导出速度慢还会占用不必要的显存。我的习惯写法是model.eval() with torch.no_grad(): torch.onnx.export( model, dummy_input, model.onnx, opset_version17, input_names[input], output_names[output], dynamic_axesdynamic_axes )另外如果模型是分布式训练出来的记住在导出前把model.module剥开转成单卡模型再导出。我遇过一次module.xxx前缀出现在 ONNX 节点名里的情况排查了半天浪费了一下午。2. ONNX 导出的核心细节与参数解析真正执行导出这一步说实话代码不复杂复杂的是每个参数背后的取舍和后果。我慢慢拆开讲一个参数一个参数过把这些细节都讲明白后面你碰到问题就知道去哪查了。2.1 opset_version 的选择逻辑opset_version是导出时最容易被忽视、却也最影响后续兼容性的参数。ONNX 的算子集opsets是版本化的不同版本支持不同的算子集合。比如Split算子在 opset 13 之前要求split参数作为 attribute到了 13 之后变成了动态输入写法完全变了又比如ReduceSum在 opset 13 之前也是 attribute13 之后变成输入。opset 版本太高导出的模型文件里会用到新版本算子老版本的推理引擎认不出来opset 版本太低一些新模型结构根本导不出来或者导出了但算子表达效率极低。我的经验是opset 13 是当前兼容性最优的版本如果你要部署到老平台比如一些边缘设备上的 ONNX Runtime 1.7 以下建议降到 11如果只是服务端部署用 17 或 18 都没问题。这次我选了 17因为后续可能要走 TensorRT太老版本的算子会导致 TRT 转换失败。还有一个细节opset_version设高不代表模型就一定能跑得更快。ONNX 只是中间表示推理速度取决于后端引擎的优化能力。高版本算子只是为了表达能力更强、减少图转换的换算损耗不要被“版本越高越快”这种错觉带偏。2.2 input_names / output_names 的规范命名这件事我在准备阶段提过但导出时还要再强调一遍因为很多人把名字写好了却发现推理时取不到输出。原因是如果模型输出是一个字典或者一个列表torch.onnx.export默认只会把第一个输出导出为 graph output其余的输出直接丢失。这时候必须显式指定output_names列表列表长度和模型实际输出数量一致。假设模型输出是(boxes, scores, classes)三元组导出代码就要写成output_names[boxes, scores, classes]这里有个隐藏的坑output_names的数量和顺序必须和模型 forward 返回的顺序一致。如果你返回的是一个 dict比如{boxes: ..., scores: ...}直接导出会报错ONNX 不支持 dict 类型的输出必须先手动拆开成 tuple 再导出。这也是 ONNX 导出前需要“模型改造”的常见场景。2.3 do_constant_folding 的取舍do_constant_folding默认是 True也就是把一些可以静态计算的节点在导出时就折叠成常量。比如模型里某个权重经过了转置、加 bias、再乘 scale这些操作在推理时输入不变所以导出时可以直接把结果算好写进 ONNX 的初始权重里减少推理时的计算量。但 constant folding 也有翻车的时候。最常见的情况是折叠后的 ONNX 模型在某个特殊输入下输出和 PyTorch 不一致。这是因为某些随机性操作如 dropout 的推理模式其实是无操作但某些自定义 op 里带了随机数被错误折叠了。我的建议是如果导出后验证发现数值不一致先试试把do_constant_foldingFalse看问题是否消失。如果消失就是折叠引入了错误如果还在那就是模型结构或算子映射的问题。2.4 验证 ONNX 模型的那几步导出完不等于工作完哪怕代码跑通没有报错也必须做验证。我在验证阶段一般做三件事。第一件事是用onnx.checker.check_model检查模型文件的结构完整性import onnx model onnx.load(model.onnx) onnx.checker.check_model(model)这个检查只能发现低级错误比如节点输入输出数量不匹配、tensor shape 不一致等。它不能保证模型推理结果正确所以还要有下一步。第二件事是拿同一份测试数据分别跑 PyTorch 模型和 ONNX 模型比对输出数值。这里我建议采集至少 3 组不同 shape 的输入如果是动态轴或 3 组同 shape 但不同数值的输入确保覆盖不同数据分布。第三件事是统计最大绝对误差。一般来说float32 精度下PyTorch 和 ONNX 的最大绝对误差在 1e-5 量级以内是正常的超过 1e-3 就说明算子映射有问题。如果跑的是量化后的 int8 模型误差放宽到 1e-2 也合理。我会把每组输入的误差记录下来作为模型交付报告的附件后续排查谁都能看懂。3. 实操过程从 PyTorch 到 ONNX 的完整落地前面讲了一堆理论现在我把这次实操的完整流程走一遍从导出脚本组织到验证脚本全部贴出来顺便把每步的关键抉择讲透。3.1 完整导出脚本这次导出的目标模型是一个 FPN 结构的检测头模型输入[1, 3, 640, 640]输出两个分支一个分类输出[1, 4, 8400, 80]一个回归输出[1, 4, 8400, 64]。完整导出脚本如下import torch import torch.nn as nn from model import build_model class ModelWrapper(nn.Module): 包装模型把多输出拆分进行命名 def __init__(self, model): super().__init__() self.model model def forward(self, x): cls, reg self.model(x) return cls, reg model build_model(pretrainedTrue) model.eval() # 构造包装模型保证输出是 tuple 而不是 dict wrapper ModelWrapper(model) # 生成 dummy input dummy_input torch.randn(1, 3, 640, 640) dynamic_axes { input: {0: batch}, cls_out: {0: batch}, reg_out: {0: batch}, } torch.onnx.export( wrapper, dummy_input, model_fpn.onnx, export_paramsTrue, opset_version17, do_constant_foldingTrue, input_names[input], output_names[cls_out, reg_out], dynamic_axesdynamic_axes, )这里我特意做了一层ModelWrapper包装原因有二一是原来模型 forward 返回的是一个 dictONNX 不接受 dict 输出必须拆成 tuple二是加了包装类以后模型内部的代码不用动后续更新模型版本时包装层保持不变代码解耦。关于输入端命名还有一个建议用input这种名字时注意不要和 ONNX 内部保留关键字冲突。我之前用data作为输入名在 onnxruntime 里取输入时没什么问题但转 TensorRT 时报了一个“data is a reserved name”的错误。后来统一改成input问题消失。3.2 导出后的文件信息检查导出完成后除了跑 checker我还习惯打印模型文件的输入输出信息确认 shape 元信息符合预期import onnx model onnx.load(model_fpn.onnx) graph model.graph for inp in graph.input: print(Input:, inp.name, inp.type.tensor_type.shape) for out in graph.output: print(Output:, out.name, out.type.tensor_type.shape)这一步能快速发现一个常见问题动态轴没有生效。比如你明明设置了dynamic_axes{0: batch}但打印出来输入第一维还是写死的 1那就说明动态轴配置没写对或者模型结构里的某个 reshape 把动态维给“固定”住了。实测中我发现某些模型结构里存在view或reshape操作会把动态维拆成[1, -1, 80]这种写法到了 ONNX 图里就变成固定数字。这种问题排查起来很费劲通常得在模型代码里把view改成reshape或者在view前引入-1的 shape 重塑。这一部分属于模型结构层面的限制最好在模型设计阶段就避免使用过多依赖固定 shape 的操作。3.3 精度校准与输出比对脚本导出后最重要的验证环节是把 PyTorch 模型和 ONNX 模型的结果做比对。这一步绝对不能省。我之前接过一个同事导出的模型导出时没有报错但部署后推理结果明显异常原因就是模型的输出经过了 sigmoid 后又被 ONNX 的 constant folding 做了错误合并。如果当时做了精度比对这个问题在导出当天就能发现。我的比对脚本思路如下import onnxruntime as ort import numpy as np import torch # 准备测试数据固定随机种保证可复现 np.random.seed(0) dummy_input np.random.randn(1, 3, 640, 640).astype(np.float32) # PyTorch 推理 with torch.no_grad(): torch_out wrapper(torch.from_numpy(dummy_input)) torch_out [o.numpy() for o in torch_out] # ONNX 推理 sess ort.InferenceSession(model_fpn.onnx, providers[CPUExecutionProvider]) onnx_out sess.run(None, {input: dummy_input}) # 逐输出比对最大绝对误差 for i, (t, o) in enumerate(zip(torch_out, onnx_out)): diff np.abs(t - o).max() print(foutput[{i}] max abs diff: {diff:.6f}) assert diff 1e-4, foutput[{i}] diff too large阈值这里我用了 1e-4偏严格。如果模型里有 softmax、layernorm 这类对浮点误差敏感的算子最大绝对误差可能会到 1e-3。但如果是卷积、ReLU 这类稳定算子1e-5 以内才算正常。如果超过 1e-2基本可以断定算子映射出了问题需要回到模型结构层面排查。一个容易踩的细节onnxruntime.InferenceSession的providers参数。默认情况下它会优先使用 CUDA provider但某些环境 CUDA 的算子实现有 bug会用一种奇怪的结果让误差变大。我建议第一次比对时全部用 CPU provider排除 GPU 端的误差干扰确认模型结构没问题后再切 GPU。4. 常见问题与排坑实录做模型导出 ONNX 的经历基本就是踩坑记录史。我把这次遇到的典型问题整理成一张速查表顺便把解决问题时的思路也写出来希望能帮你省下不少排查时间。4.1 常见错误速查表错误现象可能原因解决方案导出时报Missing op/Unsupported operatoropset 版本过低模型里的算子不在该版本支持列表调高 opset_version或替换自定义算子为标准算子导出时报Failed to export模型中有数据依赖输出或控制流分支用torch.jit.script改造成 script 模型或者重构模型消除控制流导出成功但推理结果差异大constant folding 错误折叠了带随机性的算子设置do_constant_foldingFalse重新导出动态轴配置了但是不生效模型里的 reshape/view 操作固定了 shape把 view 改为 reshape或检查动态轴对应维度是否确实可变化导出的模型文件巨大好几百MB权重没有量化或者模型里包含了不必要的辅助输出检查是否导出了几个辅助头去掉不用的输出后续再做 int8 量化ONNX Runtime 报Invalid Graph图结构里出现了孤立节点或非法 shape用onnx.checker.check_model查找具体报错节点对照修改模型结构4.2 动态 shape 场景下的坑动态 shape 是我这次遇到的最大坑没有之一。我一开始导出的模型是固定输入[1, 3, 640, 640]在服务端推理时一切正常。后来要接一个在线 API用户上传的图片分辨率不一致我就在导出时加了dynamic_axes将高度和宽度都设成动态。导出没有任何报错模型文件大小也没变化于是我直接把新模型丢到线上灰度。结果发现推理耗时比之前涨了足足 40%而且偶发 OOM。查日志才知道动态 shape 时 onnxruntime 需要对每个 batch 的输入做动态内存分配而我原本在静态 shape 下的显存池策略全部失效。更糟的是某些 resize 算子在动态 shape 下会触发 CPU 回退导致 GPU 推理变成了 CPU-GPU 混合推理性能直接崩塌。这个问题的本质在于动态轴的功能是给“输入尺寸变化”提供兼容但代价是不再享受后端引擎针对静态 shape 的预设优化。如果服务端能统一输入尺寸比如固定长边缩放、letterbox 填充就不要用动态轴。我最后的解决方案是服务端统一做 letterbox 预处理把所有图片缩放到 640x640模型保持静态 shape性能损失完全消失推理耗时回到原来的水平。4.3 量化前置操作的要点虽然当前 Phase A 还没做到量化但 ONNX 导出时就要为后续的 int8 量化预留条件。不少人问“.onnx量化int8怎么操作”我只能说导出时就做好这几件事后面量化才顺第一导出时避免使用大量的Add算子做 bias 融合。量化时算子融合是个大工程能导出的图越简洁量化工具越容易跑通。我见过一个模型原本一个卷积后面跟着 bias、BN、ReLU导出后变成了 Conv Add Mul Clip 一串节点量化时这些节点全部要重新融合工作量翻了好几倍。第二如果模型里有Gather算子常见于 embedding 层导出时确认它没有被折叠成常量。Gather 在 int8 量化里有特殊处理路径折叠成常量后反而会破坏量化流程。第三保留一个导出的原始浮点模型副本。量化过程大概率需要反复调整校准数据每次重新导出不仅慢而且可能引入新的结构变化。有一个固定的浮点版本作为基准比较量化前后误差时会方便很多。4.4 ONNX 转其他平台的注意点基于热词搜索看到的趋势最近很多人也在做 ONNX 到其他平台格式的转换比如.onnx转.kmodel、TensorRT engine、OpenVINO IR 等。我的体会是ONNX 只是一个中间表示不是终点。不同的目标平台对不同算子的支持度差异很大。比如部分边缘 NPU 对Resize算子的支持不完整导出 ONNX 时如果用了Resize的coordinate_transformation_mode参数且设为align_cornersTrue到了 NPU 上可能直接不支持或精度异常。这种问题往往要到转换阶段才能发现再回改 PyTorch 模型代价很大。提前的做法是在模型设计阶段就按“最终部署平台的算子白名单”约束网络结构。比如目标平台不支持某个自定义算子的 int8 量化那就用多个标准算子组合替代比如平台对某些算子的动态 shape 不支持那就避免过度使用动态轴。这些约束写进模型设计文档能在导出阶段省掉很多“先导出再排查”的无用功。5. 性能验证与交付清单导出工作做完关键一步是性能验证和交付物整理。很多团队把“能导出”等同于“完成”结果到了部署阶段各种不兼容返工率极高。5.1 性能验证的维度ONNX 导出后的性能验证我分两个维度看一是准确性二是推理效率。准确性验证除了前面提到的 PyTorch vs ONNX 输出比对我还会再跑一遍端到端的指标验证。比如一个检测模型直接用 COCO mAP 或者自定义测试集上的召回率/准确率做对比确认导出后的模型在真实任务上的指标没有掉。这一步很容易被省略但如果前面的算子映射埋了雷单靠最大绝对误差是发现不了的——毕竟数值差异小不等于任务指标一致。推理效率验证我会对比 PyTorch 原模型和 ONNX 模型在同一份测试数据上的单次推理耗时、峰值显存占用。注意这里的耗时一定要在排除 warmup 之后统计别拿第一次推理的数据。onnxruntime 在第一次推理时会初始化 session 和内存池耗时比稳定值高好几倍这是正常现象不算性能问题。效率验证的代码片段大致如下import time import onnxruntime as ort import numpy as np sess ort.InferenceSession(model_fpn.onnx, providers[CUDAExecutionProvider]) input_data np.random.randn(1, 3, 640, 640).astype(np.float32) # 预热 for _ in range(10): sess.run(None, {input: input_data}) # 正式计时 times [] for _ in range(50): t0 time.perf_counter() sess.run(None, {input: input_data}) times.append(time.perf_counter() - t0) print(favg: {np.mean(times[10:]) * 1000:.2f} ms)这里times[10:]的用意是去掉前 10 次因为缓存、内存池初始化带来的异常值。实际测试中这个数值一般稳定在均值上下 2% 的范围内如果抖动过大就要检查是不是显存碎片化太严重。5.2 交付清单与后续维护一个完整的 ONNX 导出交付物不能只是一个.onnx文件。我会在每次导出模型时同步输出以下内容交付物说明model.onnx标准 ONNX 模型文件model.onnx.prototxt使用onnx库反序列化为可读文本便于评审图结构导出脚本与验证脚本保证模型可复现、可回归测试输出数值比对报告记录每组测试输入的 PyTorch/ONNX 输出最大绝对误差性能测试报告记录推理耗时、显存占用、并发吞吐等指标版本与 commit 记录记录对应的 PyTorch 模型版本、训练代码版本、ONNX 导出代码版本这些交付物里最容易缺的是最后一个“版本对应关系”。我遇到过不止一次模型 A 是某个 PyTorch 权重导出的但到了部署阶段部署方拿着一个旧版 ONNX 模型跑发现和新版 PyTorch 权重的结果对不上。查了半天才发现ONNX 文件本身没有记录源权重的 commit hash。后来我要求每次导出都在 ONNX 模型的 metadata 里写入版本信息做法很简单onnx_model.metadata_props.add(keysource_model_version, valuev1.2.0) onnx_model.metadata_props.add(keyexport_timestamp, value2025-01-15T10:30:00) onnx.save(onnx_model, model_fpn.onnx)这个 metadata 属性可以附加在 ONNX 文件头不影响推理但后续追溯版本时能省掉大量沟通成本。这个习惯我建议所有做模型导出的同学都养成。6. 个人经验总结与后续规划Phase A 第 3 步的 ONNX 导出到这里基本完成了。从准备、导出、验证到交付每一步都有很多细节但核心思路其实一句话就能说清ONNX 导出的目标不是“成功跑出文件”而是“生成一个符合部署需求的、可验证的、可维护的中间产物”。我个人在这次实操中最大的体会是导出环节的很多“坑”根源不在导出本身而在模型设计阶段。比如动态分支会导致 trace 出问题如果不从模型结构层面消除数据依赖分支再怎么调导出参数都是治标不治本。所以如果你现在正处在模型设计阶段我建议提前把“能否顺利导出 ONNX”作为一个设计约束加进去而不是等到导出时才来补救。这次导出的浮点 ONNX 模型后续已经排进了两个方向一是做 int8 量化目标是模型体积压缩到原来的四分之一推理速度提升 2 倍以上二是尝试直接在 ONNX 基础上做 CPU 推理优化看看能否完全脱离 GPU 依赖降低线上的部署成本。如果你正在做类似的 ONNX 导出有一件事我特别想再强调一下导出后一定第一时间做精度校验。我自己好几次因为赶进度导出完没做校验就直接丢给下游结果下游跑了一天发现结果不对回来重新排查浪费的时间和精力远超当场跑一次比对脚本。就这几分钟的功夫真的值得花。
1. 导出前的准备工作与方案选型做模型导出 ONNX 这事乍一看好像就是把torch.onnx.export跑一遍的事但实际上十个模型九个坑前面准备不到位后面导出、验证、上板子全都是重新返工的节奏。我这次做的是 Phase A 第 3 步目标是把训练好的 PyTorch 模型导成 ONNX为后续的量化、推理加速和端侧部署铺路。整个流水线里ONNX 导出算是最关键的中间环节——它卡在“训练”和“部署”之间导不好后面所有步骤都是空中楼阁。先说准备工作。这个阶段看起来是“准备”实际上直接决定了后面能不能少折腾。第一件事是锁版本。PyTorch 版本、ONNX 版本、onnxruntime 版本这三者必须匹配。我见过不少网友在群里截图报错Failed to export an ONNX attribute onnx::Gather这类问题翻来覆去看了半天最后发现是 torch 1.8 和最新版 ONNX 不兼容。我的建议是PyTorch 用 1.13 以上最好是 2.xONNX 库用 1.13 以上onnxruntime 用最新稳定版。这三个版本组合基本能覆盖绝大多数模型的导出需求。如果你用的是老框架遇到算符缺失的问题第一反应不是改代码而是看看版本是不是太旧了。第二件事是确定导出协议。ONNX 导出有两套逻辑一套是 trace追踪一套是 script脚本化。PyTorch 默认走的是 trace 机制也就是给定一组示例输入把计算图里的实际执行路径记录下来。这意味着如果你的模型里有if data.sum() 0:这类依赖数据值的分支trace 只会记录满足当前输入条件的那一条路径另一个分支永远不会出现在 ONNX 图里。我之前导出过一个带 attention mask 的模型就是因为 mask 恰好全为 1trace 阶段把 mask 分支直接“优化”掉了导出来的模型在真实场景下结果完全错误。所以使用 trace 之前要么保证模型结构是静态的要么用torch.jit.script把动态分支显式 script 化要么干脆在设计模型时避免数据依赖的分支。第三件事是输入输出的命名规划。这个容易被忽略但到了部署阶段不管是 onnxruntime 还是 TensorRT都要按名字取输入输出。如果导出时没有显式指定input_names和output_namesONNX 会自动生成input.1、output.1这种名字你后面的推理代码就会变成魔法字符串满天飞。我习惯在导出前就把输入命名为input、输出命名为output如果模型有多个输入就在名字里加后缀比如input_ids、attention_mask一目了然。1.1 动态轴与静态轴怎么选动态轴dynamic axes和静态轴static axes的选择是我在导出前想得最多的事。所谓动态轴就是允许模型在推理时接收不同尺寸的输入。比如一个目标检测模型训练时输入图片是 640x640但部署时用户上传的图片五花八门这时候如果不做动态轴ONNX 模型就固定接收 640x640 的输入其他尺寸一律报错或强制 resize灵活性大打折扣。但动态轴不是免费午餐。动态轴会在模型图里引入额外的 Shape 计算和 Gather 算符这会带来不小的性能开销尤其是在 TensorRT 这类对静态 shape 做了深度优化的推理引擎上动态 shape 往往导致 TensorRT 回退到 fallback 模式推理速度直接掉一大截。所以我的建议是能静态就静态必须动态再动态。如果是服务于固定输入尺寸的推理服务比如视频流分析、固定分辨率检测静态轴是最稳妥的选择如果是给开放平台用的通用模型那就老老实实配dynamic_axes。举个例子我这次导出的模型输入是[1, 3, 640, 640]我用了动态轴只作用于 batch 维度和高度宽度维度代码长这样dynamic_axes { input: {0: batch, 2: height, 3: width}, output: {0: batch, 2: height, 3: width} }注意这里我没有把通道维度设为动态。因为模型结构里的卷积层权重是[64, 3, 3, 3]这种固定 shape通道数变化会引起卷积权重不匹配动态通道没有任何实际意义。类似的教训是设定动态轴时只设定确实会变化的维度不要想着“全都动态”那样会引入大量无用的 Shape 操作。1.2 模型锁定的重要性还有一个极其容易踩的坑导出前模型必须处于 eval 模式并且锁住参数。很多人在训练完直接调用 export忘记model.eval(),结果导出来的 ONNX 图里带着 dropout 层、batchnorm 层的训练行为推理时表现为输出随机抖动、精度不稳定。而且如果在导出前没有torch.no_grad()包裹模型的前向过程还会计算梯度不仅导出速度慢还会占用不必要的显存。我的习惯写法是model.eval() with torch.no_grad(): torch.onnx.export( model, dummy_input, model.onnx, opset_version17, input_names[input], output_names[output], dynamic_axesdynamic_axes )另外如果模型是分布式训练出来的记住在导出前把model.module剥开转成单卡模型再导出。我遇过一次module.xxx前缀出现在 ONNX 节点名里的情况排查了半天浪费了一下午。2. ONNX 导出的核心细节与参数解析真正执行导出这一步说实话代码不复杂复杂的是每个参数背后的取舍和后果。我慢慢拆开讲一个参数一个参数过把这些细节都讲明白后面你碰到问题就知道去哪查了。2.1 opset_version 的选择逻辑opset_version是导出时最容易被忽视、却也最影响后续兼容性的参数。ONNX 的算子集opsets是版本化的不同版本支持不同的算子集合。比如Split算子在 opset 13 之前要求split参数作为 attribute到了 13 之后变成了动态输入写法完全变了又比如ReduceSum在 opset 13 之前也是 attribute13 之后变成输入。opset 版本太高导出的模型文件里会用到新版本算子老版本的推理引擎认不出来opset 版本太低一些新模型结构根本导不出来或者导出了但算子表达效率极低。我的经验是opset 13 是当前兼容性最优的版本如果你要部署到老平台比如一些边缘设备上的 ONNX Runtime 1.7 以下建议降到 11如果只是服务端部署用 17 或 18 都没问题。这次我选了 17因为后续可能要走 TensorRT太老版本的算子会导致 TRT 转换失败。还有一个细节opset_version设高不代表模型就一定能跑得更快。ONNX 只是中间表示推理速度取决于后端引擎的优化能力。高版本算子只是为了表达能力更强、减少图转换的换算损耗不要被“版本越高越快”这种错觉带偏。2.2 input_names / output_names 的规范命名这件事我在准备阶段提过但导出时还要再强调一遍因为很多人把名字写好了却发现推理时取不到输出。原因是如果模型输出是一个字典或者一个列表torch.onnx.export默认只会把第一个输出导出为 graph output其余的输出直接丢失。这时候必须显式指定output_names列表列表长度和模型实际输出数量一致。假设模型输出是(boxes, scores, classes)三元组导出代码就要写成output_names[boxes, scores, classes]这里有个隐藏的坑output_names的数量和顺序必须和模型 forward 返回的顺序一致。如果你返回的是一个 dict比如{boxes: ..., scores: ...}直接导出会报错ONNX 不支持 dict 类型的输出必须先手动拆开成 tuple 再导出。这也是 ONNX 导出前需要“模型改造”的常见场景。2.3 do_constant_folding 的取舍do_constant_folding默认是 True也就是把一些可以静态计算的节点在导出时就折叠成常量。比如模型里某个权重经过了转置、加 bias、再乘 scale这些操作在推理时输入不变所以导出时可以直接把结果算好写进 ONNX 的初始权重里减少推理时的计算量。但 constant folding 也有翻车的时候。最常见的情况是折叠后的 ONNX 模型在某个特殊输入下输出和 PyTorch 不一致。这是因为某些随机性操作如 dropout 的推理模式其实是无操作但某些自定义 op 里带了随机数被错误折叠了。我的建议是如果导出后验证发现数值不一致先试试把do_constant_foldingFalse看问题是否消失。如果消失就是折叠引入了错误如果还在那就是模型结构或算子映射的问题。2.4 验证 ONNX 模型的那几步导出完不等于工作完哪怕代码跑通没有报错也必须做验证。我在验证阶段一般做三件事。第一件事是用onnx.checker.check_model检查模型文件的结构完整性import onnx model onnx.load(model.onnx) onnx.checker.check_model(model)这个检查只能发现低级错误比如节点输入输出数量不匹配、tensor shape 不一致等。它不能保证模型推理结果正确所以还要有下一步。第二件事是拿同一份测试数据分别跑 PyTorch 模型和 ONNX 模型比对输出数值。这里我建议采集至少 3 组不同 shape 的输入如果是动态轴或 3 组同 shape 但不同数值的输入确保覆盖不同数据分布。第三件事是统计最大绝对误差。一般来说float32 精度下PyTorch 和 ONNX 的最大绝对误差在 1e-5 量级以内是正常的超过 1e-3 就说明算子映射有问题。如果跑的是量化后的 int8 模型误差放宽到 1e-2 也合理。我会把每组输入的误差记录下来作为模型交付报告的附件后续排查谁都能看懂。3. 实操过程从 PyTorch 到 ONNX 的完整落地前面讲了一堆理论现在我把这次实操的完整流程走一遍从导出脚本组织到验证脚本全部贴出来顺便把每步的关键抉择讲透。3.1 完整导出脚本这次导出的目标模型是一个 FPN 结构的检测头模型输入[1, 3, 640, 640]输出两个分支一个分类输出[1, 4, 8400, 80]一个回归输出[1, 4, 8400, 64]。完整导出脚本如下import torch import torch.nn as nn from model import build_model class ModelWrapper(nn.Module): 包装模型把多输出拆分进行命名 def __init__(self, model): super().__init__() self.model model def forward(self, x): cls, reg self.model(x) return cls, reg model build_model(pretrainedTrue) model.eval() # 构造包装模型保证输出是 tuple 而不是 dict wrapper ModelWrapper(model) # 生成 dummy input dummy_input torch.randn(1, 3, 640, 640) dynamic_axes { input: {0: batch}, cls_out: {0: batch}, reg_out: {0: batch}, } torch.onnx.export( wrapper, dummy_input, model_fpn.onnx, export_paramsTrue, opset_version17, do_constant_foldingTrue, input_names[input], output_names[cls_out, reg_out], dynamic_axesdynamic_axes, )这里我特意做了一层ModelWrapper包装原因有二一是原来模型 forward 返回的是一个 dictONNX 不接受 dict 输出必须拆成 tuple二是加了包装类以后模型内部的代码不用动后续更新模型版本时包装层保持不变代码解耦。关于输入端命名还有一个建议用input这种名字时注意不要和 ONNX 内部保留关键字冲突。我之前用data作为输入名在 onnxruntime 里取输入时没什么问题但转 TensorRT 时报了一个“data is a reserved name”的错误。后来统一改成input问题消失。3.2 导出后的文件信息检查导出完成后除了跑 checker我还习惯打印模型文件的输入输出信息确认 shape 元信息符合预期import onnx model onnx.load(model_fpn.onnx) graph model.graph for inp in graph.input: print(Input:, inp.name, inp.type.tensor_type.shape) for out in graph.output: print(Output:, out.name, out.type.tensor_type.shape)这一步能快速发现一个常见问题动态轴没有生效。比如你明明设置了dynamic_axes{0: batch}但打印出来输入第一维还是写死的 1那就说明动态轴配置没写对或者模型结构里的某个 reshape 把动态维给“固定”住了。实测中我发现某些模型结构里存在view或reshape操作会把动态维拆成[1, -1, 80]这种写法到了 ONNX 图里就变成固定数字。这种问题排查起来很费劲通常得在模型代码里把view改成reshape或者在view前引入-1的 shape 重塑。这一部分属于模型结构层面的限制最好在模型设计阶段就避免使用过多依赖固定 shape 的操作。3.3 精度校准与输出比对脚本导出后最重要的验证环节是把 PyTorch 模型和 ONNX 模型的结果做比对。这一步绝对不能省。我之前接过一个同事导出的模型导出时没有报错但部署后推理结果明显异常原因就是模型的输出经过了 sigmoid 后又被 ONNX 的 constant folding 做了错误合并。如果当时做了精度比对这个问题在导出当天就能发现。我的比对脚本思路如下import onnxruntime as ort import numpy as np import torch # 准备测试数据固定随机种保证可复现 np.random.seed(0) dummy_input np.random.randn(1, 3, 640, 640).astype(np.float32) # PyTorch 推理 with torch.no_grad(): torch_out wrapper(torch.from_numpy(dummy_input)) torch_out [o.numpy() for o in torch_out] # ONNX 推理 sess ort.InferenceSession(model_fpn.onnx, providers[CPUExecutionProvider]) onnx_out sess.run(None, {input: dummy_input}) # 逐输出比对最大绝对误差 for i, (t, o) in enumerate(zip(torch_out, onnx_out)): diff np.abs(t - o).max() print(foutput[{i}] max abs diff: {diff:.6f}) assert diff 1e-4, foutput[{i}] diff too large阈值这里我用了 1e-4偏严格。如果模型里有 softmax、layernorm 这类对浮点误差敏感的算子最大绝对误差可能会到 1e-3。但如果是卷积、ReLU 这类稳定算子1e-5 以内才算正常。如果超过 1e-2基本可以断定算子映射出了问题需要回到模型结构层面排查。一个容易踩的细节onnxruntime.InferenceSession的providers参数。默认情况下它会优先使用 CUDA provider但某些环境 CUDA 的算子实现有 bug会用一种奇怪的结果让误差变大。我建议第一次比对时全部用 CPU provider排除 GPU 端的误差干扰确认模型结构没问题后再切 GPU。4. 常见问题与排坑实录做模型导出 ONNX 的经历基本就是踩坑记录史。我把这次遇到的典型问题整理成一张速查表顺便把解决问题时的思路也写出来希望能帮你省下不少排查时间。4.1 常见错误速查表错误现象可能原因解决方案导出时报Missing op/Unsupported operatoropset 版本过低模型里的算子不在该版本支持列表调高 opset_version或替换自定义算子为标准算子导出时报Failed to export模型中有数据依赖输出或控制流分支用torch.jit.script改造成 script 模型或者重构模型消除控制流导出成功但推理结果差异大constant folding 错误折叠了带随机性的算子设置do_constant_foldingFalse重新导出动态轴配置了但是不生效模型里的 reshape/view 操作固定了 shape把 view 改为 reshape或检查动态轴对应维度是否确实可变化导出的模型文件巨大好几百MB权重没有量化或者模型里包含了不必要的辅助输出检查是否导出了几个辅助头去掉不用的输出后续再做 int8 量化ONNX Runtime 报Invalid Graph图结构里出现了孤立节点或非法 shape用onnx.checker.check_model查找具体报错节点对照修改模型结构4.2 动态 shape 场景下的坑动态 shape 是我这次遇到的最大坑没有之一。我一开始导出的模型是固定输入[1, 3, 640, 640]在服务端推理时一切正常。后来要接一个在线 API用户上传的图片分辨率不一致我就在导出时加了dynamic_axes将高度和宽度都设成动态。导出没有任何报错模型文件大小也没变化于是我直接把新模型丢到线上灰度。结果发现推理耗时比之前涨了足足 40%而且偶发 OOM。查日志才知道动态 shape 时 onnxruntime 需要对每个 batch 的输入做动态内存分配而我原本在静态 shape 下的显存池策略全部失效。更糟的是某些 resize 算子在动态 shape 下会触发 CPU 回退导致 GPU 推理变成了 CPU-GPU 混合推理性能直接崩塌。这个问题的本质在于动态轴的功能是给“输入尺寸变化”提供兼容但代价是不再享受后端引擎针对静态 shape 的预设优化。如果服务端能统一输入尺寸比如固定长边缩放、letterbox 填充就不要用动态轴。我最后的解决方案是服务端统一做 letterbox 预处理把所有图片缩放到 640x640模型保持静态 shape性能损失完全消失推理耗时回到原来的水平。4.3 量化前置操作的要点虽然当前 Phase A 还没做到量化但 ONNX 导出时就要为后续的 int8 量化预留条件。不少人问“.onnx量化int8怎么操作”我只能说导出时就做好这几件事后面量化才顺第一导出时避免使用大量的Add算子做 bias 融合。量化时算子融合是个大工程能导出的图越简洁量化工具越容易跑通。我见过一个模型原本一个卷积后面跟着 bias、BN、ReLU导出后变成了 Conv Add Mul Clip 一串节点量化时这些节点全部要重新融合工作量翻了好几倍。第二如果模型里有Gather算子常见于 embedding 层导出时确认它没有被折叠成常量。Gather 在 int8 量化里有特殊处理路径折叠成常量后反而会破坏量化流程。第三保留一个导出的原始浮点模型副本。量化过程大概率需要反复调整校准数据每次重新导出不仅慢而且可能引入新的结构变化。有一个固定的浮点版本作为基准比较量化前后误差时会方便很多。4.4 ONNX 转其他平台的注意点基于热词搜索看到的趋势最近很多人也在做 ONNX 到其他平台格式的转换比如.onnx转.kmodel、TensorRT engine、OpenVINO IR 等。我的体会是ONNX 只是一个中间表示不是终点。不同的目标平台对不同算子的支持度差异很大。比如部分边缘 NPU 对Resize算子的支持不完整导出 ONNX 时如果用了Resize的coordinate_transformation_mode参数且设为align_cornersTrue到了 NPU 上可能直接不支持或精度异常。这种问题往往要到转换阶段才能发现再回改 PyTorch 模型代价很大。提前的做法是在模型设计阶段就按“最终部署平台的算子白名单”约束网络结构。比如目标平台不支持某个自定义算子的 int8 量化那就用多个标准算子组合替代比如平台对某些算子的动态 shape 不支持那就避免过度使用动态轴。这些约束写进模型设计文档能在导出阶段省掉很多“先导出再排查”的无用功。5. 性能验证与交付清单导出工作做完关键一步是性能验证和交付物整理。很多团队把“能导出”等同于“完成”结果到了部署阶段各种不兼容返工率极高。5.1 性能验证的维度ONNX 导出后的性能验证我分两个维度看一是准确性二是推理效率。准确性验证除了前面提到的 PyTorch vs ONNX 输出比对我还会再跑一遍端到端的指标验证。比如一个检测模型直接用 COCO mAP 或者自定义测试集上的召回率/准确率做对比确认导出后的模型在真实任务上的指标没有掉。这一步很容易被省略但如果前面的算子映射埋了雷单靠最大绝对误差是发现不了的——毕竟数值差异小不等于任务指标一致。推理效率验证我会对比 PyTorch 原模型和 ONNX 模型在同一份测试数据上的单次推理耗时、峰值显存占用。注意这里的耗时一定要在排除 warmup 之后统计别拿第一次推理的数据。onnxruntime 在第一次推理时会初始化 session 和内存池耗时比稳定值高好几倍这是正常现象不算性能问题。效率验证的代码片段大致如下import time import onnxruntime as ort import numpy as np sess ort.InferenceSession(model_fpn.onnx, providers[CUDAExecutionProvider]) input_data np.random.randn(1, 3, 640, 640).astype(np.float32) # 预热 for _ in range(10): sess.run(None, {input: input_data}) # 正式计时 times [] for _ in range(50): t0 time.perf_counter() sess.run(None, {input: input_data}) times.append(time.perf_counter() - t0) print(favg: {np.mean(times[10:]) * 1000:.2f} ms)这里times[10:]的用意是去掉前 10 次因为缓存、内存池初始化带来的异常值。实际测试中这个数值一般稳定在均值上下 2% 的范围内如果抖动过大就要检查是不是显存碎片化太严重。5.2 交付清单与后续维护一个完整的 ONNX 导出交付物不能只是一个.onnx文件。我会在每次导出模型时同步输出以下内容交付物说明model.onnx标准 ONNX 模型文件model.onnx.prototxt使用onnx库反序列化为可读文本便于评审图结构导出脚本与验证脚本保证模型可复现、可回归测试输出数值比对报告记录每组测试输入的 PyTorch/ONNX 输出最大绝对误差性能测试报告记录推理耗时、显存占用、并发吞吐等指标版本与 commit 记录记录对应的 PyTorch 模型版本、训练代码版本、ONNX 导出代码版本这些交付物里最容易缺的是最后一个“版本对应关系”。我遇到过不止一次模型 A 是某个 PyTorch 权重导出的但到了部署阶段部署方拿着一个旧版 ONNX 模型跑发现和新版 PyTorch 权重的结果对不上。查了半天才发现ONNX 文件本身没有记录源权重的 commit hash。后来我要求每次导出都在 ONNX 模型的 metadata 里写入版本信息做法很简单onnx_model.metadata_props.add(keysource_model_version, valuev1.2.0) onnx_model.metadata_props.add(keyexport_timestamp, value2025-01-15T10:30:00) onnx.save(onnx_model, model_fpn.onnx)这个 metadata 属性可以附加在 ONNX 文件头不影响推理但后续追溯版本时能省掉大量沟通成本。这个习惯我建议所有做模型导出的同学都养成。6. 个人经验总结与后续规划Phase A 第 3 步的 ONNX 导出到这里基本完成了。从准备、导出、验证到交付每一步都有很多细节但核心思路其实一句话就能说清ONNX 导出的目标不是“成功跑出文件”而是“生成一个符合部署需求的、可验证的、可维护的中间产物”。我个人在这次实操中最大的体会是导出环节的很多“坑”根源不在导出本身而在模型设计阶段。比如动态分支会导致 trace 出问题如果不从模型结构层面消除数据依赖分支再怎么调导出参数都是治标不治本。所以如果你现在正处在模型设计阶段我建议提前把“能否顺利导出 ONNX”作为一个设计约束加进去而不是等到导出时才来补救。这次导出的浮点 ONNX 模型后续已经排进了两个方向一是做 int8 量化目标是模型体积压缩到原来的四分之一推理速度提升 2 倍以上二是尝试直接在 ONNX 基础上做 CPU 推理优化看看能否完全脱离 GPU 依赖降低线上的部署成本。如果你正在做类似的 ONNX 导出有一件事我特别想再强调一下导出后一定第一时间做精度校验。我自己好几次因为赶进度导出完没做校验就直接丢给下游结果下游跑了一天发现结果不对回来重新排查浪费的时间和精力远超当场跑一次比对脚本。就这几分钟的功夫真的值得花。
