PyTorch导出ONNX全流程：动态batch、opset、数值对齐与部署

模型训练完只是半成品，能不能落到业务线上，八成卡在导出这一步。这两年我经手的推理部署项目里，十有八九要把 PyTorch 训练出来的权重转成 ONNX，再交给后端的推理引擎去跑。PyTorch导出模型为onnx这件事，表面上看就是一行torch.onnx.export的调用，真上手才知道坑密得像筛子：动态 batch 导出来变成写死 1、自定义算子直接报不支持、数值对不齐差出小数点后三位、导出的模型在本地跑得好好的到了线上就崩。这篇东西写给已经会训模型、正准备把模型往外交付的人，也写给被导出问题折腾过几轮、想把流程做扎实的人。我会把方案选型、关键参数、验证手段、排错路径和工程化落地一次讲透，尽量让你少走我走过的弯路。

1. 导出这件事到底在解决什么问题

1.1 ONNX 在推理链路里的位置

先把概念摆正。ONNX 不是模型格式的终点，它是一个中间表示层。你在 PyTorch 里定义的计算图，本质上是 Python 对象加上一堆 eager 执行的算子调用；而 ONNX 描述的是一张静态的、与框架无关的计算图，节点是标准算子，边是张量，附带形状和类型信息。训练框架负责产出这张图，推理引擎负责消费这张图，两者解耦，这就是整个生态存在的理由。

理解这一点很关键，因为它决定了你导出的目标不是"存一份权重"，而是"翻译一张图"。权重会被塞进图里变成 initializer，结构会被翻译成节点，而翻译过程中任何 PyTorch 特有、ONNX 标准里没有对应物的写法，都会变成麻烦。我见过太多人把导出当成torch.save的替代品，结果导出来的模型根本跑不动，或者跑出来的结果和原模型对不上。

从工程角度看，ONNX 带来的收益集中在三块。一是跨框架、跨语言部署，C++、C#、Java、Python 都能加载同一份文件；二是推理引擎可以做图层面的优化，算子融合、常量折叠、内存复用这些手段在 eager 模式下是没法做的；三是交付边界清晰，你把一个自包含的文件交给下游，对方不需要装整套训练框架。这三点决定了它在生产环境里的位置，也决定了你对导出质量的容忍度应该有多低。

1.2 什么样的项目真的需要导出

不是所有项目都值得走这条路。如果你的模型只在训练脚本里做评估，或者服务端本来就是 Python 环境、QPS 也不高，直接加载.pt用 PyTorch 跑反而省事。导出带来的额外成本是实实在在的：你要维护一套对齐验证，要处理算子兼容，要跟着框架版本升级重新适配。

真正需要导出的场景，我归纳成几类。第一类是端侧部署，手机、嵌入式板子、浏览器，这些环境压根装不了完整训练框架。第二类是高吞吐服务，推理引擎在同样的硬件上能压出更高的吞吐和更低的尾延迟，尤其是 batch 推理场景。第三类是异构硬件，专用加速卡通常只认标准图格式，你得先把图交出去。第四类是多团队协作，算法团队交模型、工程团队做服务，中间需要一个格式契约，ONNX 就是那个契约。

反过来说，如果你的模型里大量依赖 Python 控制流、动态数据结构、运行时才知道的 shape，那导出的代价会非常高，可能改模型结构比导出本身还费劲。这种时候要先评估：是改造模型划算，还是保留原框架直接部署划算。我个人的判断线是，如果为了让模型可导出需要改动超过三成的核心逻辑，就先别急着导，先看看业务侧的真实瓶颈在哪里。

1.3 导出前先想清楚的三件事

动手前有三个问题必须先有答案，否则后面全是返工。

第一个问题是目标运行时的算子覆盖范围。不同的推理引擎支持的 opset 版本和算子集不完全一样，有些引擎对某些算子的高版本实现支持滞后。你要先确认下游能吃哪个 opset，再倒推导出时用什么版本。这个顺序不能反，先导出再问对方支持什么，大概率要重来。

第二个问题是输入的 shape 语义。你的模型是要固定 batch 还是动态 batch，图像尺寸是固定的还是可变的，序列长度是否有上限。这些决定了dynamic_axes怎么配，也决定了后续能不能做 shape 推断优化。很多线上事故的根源就是导出时把 batch 写死成 1，压测的时候才发现根本没法批量推理。

第三个问题是数值精度要求。ONNX 默认用 float32，如果你的模型里有混合精度训练出来的权重，或者对数值敏感的操作，导出前后的误差必须量化评估。我一般会在导出前先定一个误差阈值，比如分类任务 top-1 一致率必须 100%，回归任务最大绝对误差不超过 1e-4，达不到就打回重做。

2. 方案选型：三条导出路径怎么挑

2.1 传统 torch.onnx.export 路径

这是最主流、文档最全的一条路。它基于 TorchScript 的 tracing 机制，把模型跑一遍，记录下实际执行的算子序列，再翻译成 ONNX 图。优点是成熟稳定、兼容性好，绝大多数量化、视觉、NLP 模型都能走通；缺点是 tracing 只记录"跑过的那条路径"，控制流里没走到的分支会直接丢失。

用它的时候有个前提容易被忽略：模型必须处于 eval 模式。BatchNorm 在 train 模式下会用当前 batch 的统计量，Dropout 会随机置零，trace 出来的图和推理时的行为完全不是一回事。我见过导出后精度暴跌的案例，查了半天发现是忘了model.eval()。这一行代码不是可选项，是硬性要求。

另一个隐性坑是 tracing 的输入。你传进去的 dummy input 不只是用来触发前向，它的 shape 和 dtype 会直接写进图里。如果你传的是torch.randn(1, 3, 224, 224)，那导出的图默认就是 batch=1、分辨率 224。想要动态，必须在dynamic_axes里显式声明，不能指望它自动推断。

2.2 TorchDynamo 后端路径

新版本 PyTorch 提供了基于 TorchDynamo 的导出方式，走的是字节码分析而不是 tracing，对控制流的支持明显更好。Python 层的 if/for 能在一定程度上被正确捕获，而不是简单地按 trace 到的路径固化。这对含有条件分支的模型是实打实的改善。

但这条路目前的状态是"能用但要挑场景"。它对模型写法的要求更严格，一些动态的 Python 操作、自定义容器、运行时属性访问仍然会触发图断裂，断裂的部分会回退到 Python 执行，导出结果可能包含不完整的图或者直接报错。我的建议是：新项目、模型结构干净、框架版本较新，可以优先试 Dynamo；老项目、结构复杂、依赖大量自定义逻辑，先用传统路径保底。

实际选型时我会两个都试一遍，对比导出成功率和验证误差，哪个稳用哪个。不要迷信新路径，导出这件事的评判标准只有一个：下游能不能稳定跑出正确结果。

2.3 脚本化与 tracing 的取舍

除了上面两条，torch.jit.script也是一条路，先把模型编译成 TorchScript 再导出。它对控制流的支持比 trace 好，但要求模型代码符合 TorchScript 的语法子集，很多 Python 惯用写法会直接编译失败。改造工作量往往不小。

三条路径对比下来，我的经验是这样：纯前馈结构、无复杂控制流，传统 export 最快最稳；有 if/while 但结构规整，Dynamo 值得一试；模型里全是 Python 动态逻辑，先考虑重构，别硬导。下面这张表是我在实际项目里总结的选型参考。

2.4 版本矩阵与依赖锁定

导出失败的原因里，版本不匹配能占三成。PyTorch 版本、ONNX 版本、opset 版本、推理引擎版本，四者之间有一张隐性的兼容矩阵。你本地用新框架导出的图，下游用老引擎加载，报个"unsupported op"再正常不过。

我的做法是把这套版本写进项目的依赖清单，并且和下游对齐后锁死。具体要锁的包括：训练侧框架版本、导出时用的 opset、onnx 检查工具版本、推理引擎版本。这四样任何一样变动，都要重跑一遍对齐验证。别嫌麻烦，线上因为版本漂移导致模型加载失败的排查成本，比这套流程高一个数量级。

还有个小细节：导出脚本里最好显式打印各组件版本，写进产物旁边的元信息文件。半年后你回头查一个老模型的导出参数，这份记录能救命。

3. 动手之前必须搞懂的核心概念

3.1 计算图、动态轴与 shape 语义

ONNX 图里的每个张量都有 shape，但 shape 有两种：静态的和符号化的。静态 shape 就是一个确定的数字，符号化 shape 是一个带名字的维度，运行时可以是任意值。dynamic_axes干的事就是把指定维度从静态改成符号化。

这里有个认知误区要纠正：dynamic_axes不是"允许动态"，而是"声明这一维是动态的"。你不声明，它就是一个固定值；你声明了，下游才能按变长输入来做内存规划和优化。两者的差别在部署时会被放大，尤其是 batch 维度和序列长度维度。

配置的时候要按维度索引来写，容易写错。对于一个四维输入[N, C, H, W]，索引 0 是 batch，1 是通道，2 和 3 是空间维度。如果你写成了{1: "batch"}，那声明的就是通道维动态，运行时喂不同通道数的输入，模型行为会完全错乱。这种错误不会在导出时报错，只会在推理时给你一个莫名其妙的 shape mismatch。

我的习惯是给每个动态维度起有意义的名字，batch 就叫batch，高度宽度就叫height、width，序列长度就叫seq_len。名字本身不参与计算，但它是给人和给工具看的契约，出了问题一眼能定位。

3.2 opset 版本与算子支持

opset 是 ONNX 的算子集版本，每个版本会新增算子、修改已有算子的定义。选版本的核心原则是"够用且下游支持"，不是越高越好。高版本可能引入了下游引擎还没实现的新算子，低版本又可能缺你需要的算子。

我一般的选法是：先查下游推理引擎的文档，确认它支持的最高 opset；再查你模型里用到的算子在这个 opset 下是否齐全；两者取交集。如果模型里有较新的算子，比如某些注意力相关的融合算子，可能需要更高版本，那就要推动下游升级，或者在上游做等价替换。

还有一个实操细节：导出时如果某个算子在当前 opset 下没有实现，工具通常会给出 warning 并把它降级成一个组合实现，或者直接报错。warning 不能忽略，它往往意味着图里多了一堆小算子，推理性能会打折。我遇到这种情况会去看具体是哪个算子，能换写法就换写法，换不了再考虑升 opset。

3.3 eval 模式、参数冻结与常量折叠

model.eval()这件事前面提过，这里展开说清楚它影响了什么。它影响的是所有"训练态和推理态行为不同"的层：BatchNorm 用滑动统计量而不是 batch 统计量，Dropout 直接透传，一些自定义层如果有类似的开关也会切过去。导出前不切，图里固化的就是训练态行为，结果必然错。

常量折叠是另一个容易被忽视的优化。导出时工具会把能提前算出来的子图直接算掉，比如连续的 reshape、transpose、常量加法。这个开关默认是开的，绝大多数情况应该保持开启，它能让图更小、推理更快。但如果你的模型里有依赖运行时状态的常量计算，折叠可能改变语义，这时候需要关掉它并逐段核对。

参数冻结指的是把参与训练但不参与推理的分支裁掉。比如一些模型带辅助分类头，只在训练时用，导出时应该让 forward 走推理分支。这件事没法靠参数解决，只能改代码，用if self.training之类的开关把分支隔离干净，再配合 eval 模式让 trace 只走推理路径。

4. 实操：手把手导出一个可用的 ONNX 模型

4.1 环境准备与最小可运行例子

先把环境摆好。我一般用一个干净的虚拟环境，装三个东西：对应版本的 PyTorch、onnx、onnxruntime。前两个负责导出和结构检查，第三个负责推理验证。版本按前面说的矩阵锁死，别用 latest。

下面是一个最小可运行例子，结构简单但五脏俱全，包含了卷积、归一化、激活、池化和全连接，足够演示完整流程。

import torch import torch.nn as nn class DemoNet(nn.Module): def __init__(self, num_classes=10): super().__init__() self.conv = nn.Conv2d(3, 32, kernel_size=3, padding=1) self.bn = nn.BatchNorm2d(32) self.act = nn.ReLU(inplace=True) self.pool = nn.AdaptiveAvgPool2d(1) self.fc = nn.Linear(32, num_classes) def forward(self, x): x = self.conv(x) x = self.bn(x) x = self.act(x) x = self.pool(x) x = torch.flatten(x, 1) return self.fc(x) model = DemoNet().eval() dummy_input = torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, "demo_net.onnx", input_names=["image"], output_names=["logits"], dynamic_axes={ "image": {0: "batch", 2: "height", 3: "width"}, "logits": {0: "batch"}, }, opset_version=17, do_constant_folding=True, )

几个点值得单独说。model.eval()在构造完就调用了，不能省。dummy_input的 shape 决定了默认图结构，虽然声明了动态轴，但 trace 时走的还是这个具体 shape 的路径，所以 dummy 的取值要贴近真实输入分布，避免某些分支没被走到。opset_version显式指定，不依赖默认值，默认值会随框架版本变，是隐性风险。

4.2 关键参数逐个拆解

input_names和output_names看着不起眼，实际很重要。它们决定了图里输入输出节点的名字，下游加载模型时就是靠这些名字来绑定数据的。不指定的话工具会生成input.1、output.1这种自动名，可读性差，而且不同版本生成的规则可能变。我见过上下游因为名字对不上互相扯皮的情况，所以这两个参数我从来都是显式写。

dynamic_axes的坑前面讲过，这里补充一个实操技巧：导出后用 onnx 工具把图的输入输出 shape 打出来，确认动态维度确实变成了符号名而不是数字。这一步花不了一分钟，能挡掉后面一堆问题。

do_constant_folding建议保持 True。它的作用是把能提前算的常量运算折叠掉，减小图规模。唯一需要关掉的情况是模型里有依赖运行时状态的常量计算，但这种设计本身就不推荐，遇到了应该改模型而不是关优化。

opset_version的选择逻辑前面说过，补充一点：如果你不确定该用哪个，可以先按下游支持的最高版本导一次，跑通了再往下试，找到一个稳定且够用的版本。不要一开始就追求最新，稳定性优先。

还有一个参数export_params，默认 True，表示把权重一起写进文件。除非你确实要做权重外置，否则保持默认。外置权重会多一个文件，部署时容易漏拷，我一般不用。

4.3 多输入、多输出与自定义算子

多输入模型的导出，把 dummy input 换成一个元组就行，顺序和 forward 的参数顺序一致。input_names也按同样顺序列。这里最容易出错的是顺序对不上，导出能成功，推理时数据绑错位置，结果全错还查不出来。

dummy = (torch.randn(1, 3, 224, 224), torch.randn(1, 128)) torch.onnx.export( model, dummy, "multi_input.onnx", input_names=["image", "text_feat"], output_names=["logits", "probs"], dynamic_axes={"image": {0: "batch"}, "text_feat": {0: "batch"}}, )

多输出同理，forward 返回元组，output_names按顺序对应。我建议输出名带上语义，比如logits、probs、boxes，别用out0、out1，下游对接的时候能省很多沟通。

自定义算子是最麻烦的一类。ONNX 标准里没有的算子，导出会失败或者生成一个ATen域的节点，很多推理引擎不认。处理思路有三条：一是用标准算子组合出等价实现，这是最干净的方案；二是注册自定义算子，需要下游引擎配合支持，成本高；三是导出后手工改图，用等价子图替换，适合一次性场景。

我的优先级是能组合就组合，实在不行再谈自定义。组合实现可能会带来一些性能损失，但换来的是可移植性，绝大多数场景这笔账是划算的。

4.4 大模型与显存受限场景

参数量大的模型导出时有个现实问题：trace 过程需要把模型完整加载，还要跑一次前向，显存占用不小。如果单卡放不下，可以用几个手段缓解。

一是用 meta device 或者半精度加载，导出前再转回需要的精度。二是分段导出再拼图，适合结构清晰的模型，但要处理跨段边界。三是用torch.onnx.export的device参数，让导出在指定设备上完成，新版支持这个能力。

另外提醒一句，大模型导出时间可能很长，几分钟到几十分钟都有。导出过程别用交互式环境，写成脚本跑，日志重定向到文件，出问题能回溯。

5. 导出后验证：怎么确认它没悄悄坏掉

5.1 结构检查与图可视化

导出成功不等于导出正确。第一步是做结构检查，用 onnx 自带的检查器把图扫一遍，能发现算子版本不匹配、节点连接错误、shape 推断失败这类问题。这一步很快，应该写进脚本自动执行。

import onnx model = onnx.load("demo_net.onnx") onnx.checker.check_model(model) # 打印输入输出信息，确认动态轴生效 for inp in model.graph.input: print(inp.name, [d.dim_param or d.dim_value for d in inp.type.tensor_type.shape.dim]) for out in model.graph.output: print(out.name, [d.dim_param or d.dim_value for d in out.type.tensor_type.shape.dim])

打印出来的维度里，动态轴应该显示成名字（比如batch），静态轴显示成数字。如果该动态的地方显示成了 1，说明dynamic_axes没配对，回去改。

图可视化工具能把计算图渲染出来，节点和连接一目了然。我一般用它做两件事：确认关键算子确实存在、没有被降级成奇怪的组合；确认图的输入输出节点名字符合预期。图太大时不用全看，看首尾和中间几个关键节点就够。

5.2 数值对齐：用同一批输入跑两遍

结构没问题只是及格线，数值对不对才是关键。做法很直接：准备一批真实分布的输入，分别用原模型和 ONNX 模型跑一遍，逐元素比对。

import numpy as np import onnxruntime as ort sess = ort.InferenceSession("demo_net.onnx", providers=["CPUExecutionProvider"]) input_name = sess.get_inputs()[0].name x = torch.randn(4, 3, 224, 224) with torch.no_grad(): ref = model(x).numpy() onnx_out = sess.run(None, {input_name: x.numpy()})[0] max_abs_err = np.max(np.abs(ref - onnx_out)) print("max abs error:", max_abs_err)

误差阈值要看任务。分类任务我一般看预测类别是否一致，要求 100% 一致；回归任务看最大绝对误差和相对误差，阈值按业务容忍度定。如果误差超标，先别急着改模型，先排查是不是操作模式、输入 dtype、归一化处理哪里不一致。

5.3 偏差来源排查表

数值对不齐的原因有规律可循，我把常见来源整理成一张表，按出现频率排序，遇到问题从上往下查。

逐层对比是个很有效的手段。把模型拆成若干段，每段单独导出、单独比对，能快速定位是哪一层开始出现偏差。这个方法费点时间，但比盲猜高效得多。

6. 常见问题与排查速查表

6.1 典型报错与解法

导出过程里的报错五花八门，但高频的就那么几类。

RuntimeError: Only tuples, lists and Variables are supported，一般是 forward 返回值类型不对，比如返回了字典或者自定义对象。解法是改成返回元组，需要多输出的用元组组织。

Unsupported operator，算子不支持。解法是查这个算子在当前 opset 下的实现情况，能换写法就换，换不了就升 opset 或者做等价组合。

TracerWarning，trace 过程中遇到了控制流或者动态行为，图可能不完整。warning 不能无视，要认真看是哪个操作触发的，能静态化就静态化。

Shape inference failed，shape 推断失败，通常是图里有动态维度没声明清楚。解法是补齐dynamic_axes或者检查是否有运行时才确定的 shape。

6.2 动态 shape 踩坑实录

动态 shape 是问题高发区，值得单独说。最常见的错误是声明了动态轴但推理时喂了超范围的值，比如声明了动态但实际实现里有隐含假设。有些模型内部用了 view 或者 reshape，把 batch 维和空间维混在一起算，batch 一变就崩。这种问题要在导出前用多个不同 shape 的输入测一遍，别只测一个。

第二个坑是动态轴声明了但下游没用上。有些引擎在加载时会根据图里的 shape 信息做静态优化，如果它没正确识别符号维度，优化后的图可能还是按固定 shape 跑。这种情况要在引擎侧确认它对动态维度的处理策略，必要时在加载时显式指定。

第三个坑是输出维度依赖输入。比如检测模型输出框的数量和输入分辨率相关，这种动态性在 ONNX 里表达起来比较别扭，通常的做法是固定最大输出数，用无效值填充。这个方案不优雅但实用，我见过不少项目这么做。

6.3 部署端的隐性坑

导出侧做对了，部署侧照样能翻车。我遇到过几次，问题都不在模型本身。

一是输入布局不一致。PyTorch 用 NCHW，有些推理引擎默认 NHWC，加载时如果不做转换，结果会错得离谱但不会报错。这种问题要靠对齐验证发现，不能靠肉眼。

二是预处理没对齐。训练时的归一化参数、通道顺序、resize 插值方式，任何一个和部署侧不一致，都会导致精度下降。我的做法是把预处理参数写进模型元信息，部署侧照抄，不做二次解释。

三是线程数和内存配置。推理引擎的线程数、内存池大小、图优化级别都会影响性能和稳定性，默认值不一定适合你的场景。上线前要做一轮压测，把配置调到稳定区间。

7. 工程化：把导出固化成可复现的流水线

7.1 目录结构与命名约定

零散的导出脚本早晚会变成技术债。我的习惯是给每个模型建一个独立的导出目录，里面固定放几样东西：导出脚本、依赖清单、验证脚本、产物文件、元信息文件。命名带上模型名和版本号，比如detector_v3_opset17.onnx，一眼能看出是什么、什么版本、什么 opset。

元信息文件里记录：导出时间、框架版本、opset 版本、输入输出名和 shape、动态轴配置、验证误差、负责人员。这份文件的价值在半年后体现，那时候没人记得当时怎么导的，全靠它。

7.2 性能优化与量化

导出只是第一步，后面还有优化空间。图优化层面，推理引擎通常会自动做算子融合和内存复用，你可以在加载时选择优化级别。量化层面，如果业务允许精度损失，把 float32 转成 int8 能带来显著的性能提升，但需要校准数据集，并且要做充分的精度回归。

量化的坑在于校准集的选择。校准集分布和真实数据偏差大，量化后的精度会掉得很厉害。我的做法是从真实流量里采样，覆盖主要场景，样本量按模型复杂度定，一般几百到上千张足够。量化后必须重新跑一遍对齐验证，阈值可以比 float32 放宽，但不能不设。

7.3 CI 里加一道回归关卡

导出流程最怕的是"改了模型忘了重导"，或者"重导了但没验证"。解法是把导出和验证写进持续集成，模型代码一有变更就自动跑一遍，验证不通过直接卡住。

具体做法是准备一份小的验证数据集，固定随机种子，每次导出后自动比对输出。这个流程搭起来要花点时间，但能挡掉绝大多数低级错误。我在几个项目里推过这套，跑起来之后导出相关的线上问题基本归零。

8. 一些踩坑之后才明白的事

最后分享几条个人经验，都是吃过亏换来的。

dummy input 的取值比想象中重要。它不是随便填个形状就完事，取值分布要贴近真实数据，否则某些依赖数值的分支可能走到不同路径。我遇到过一次，dummy 用全零，导出后模型对某些输入输出异常，换成随机正态就正常了。原因是模型里有基于输入统计量的逻辑，全零输入走了特殊路径。

导出脚本要能重复跑。不要用交互式环境一步步敲，写成脚本，参数化，任何一次导出都能用同一条命令复现。这听起来是基本功，但我见过太多人导出参数改了没记录，出问题只能从头试。

验证要自动化。手动比对一次两次可以，长期靠人眼不现实。把验证写进脚本，每次导出自动跑，误差超标直接报错。这套机制建立起来之后，你对导出质量的信心会完全不一样。

还有一个心态上的建议：导出失败别急着换方案。先看报错，很多问题是有明确原因的，改一行代码就能解决。我见过有人一报错就换框架、换工具链，折腾一圈回到原点。导出这件事的复杂度大部分在细节，不在选型。

这个领域后续还能往深里做的是量化感知导出和跨版本兼容矩阵的自动化维护，前者能进一步压推理成本，后者能减少版本升级带来的返工。如果手头的模型对延迟敏感，这两块值得投入。