ONNX是跨框架通用的模型交换格式，把PyTorch模型转成ONNX后可以部署到TensorRT、ONNX Runtime等多种推理引擎上运行。你会学到torch.onnx.export的用法、动态维度和动态batch处理、opset版本选择策略、常见算子不支持时的自定义export方法，以及用onnxsim做模型简化。适合做模型部署、跨平台推理、端侧落地的工程师。

岗位关联

需要学习PyTorch转ONNX，这是你的核心专业课，TVM/XLA/Triton等编译器是你的日常工具。本篇涵盖TVM/XLA/Triton、算子融合、图优化，帮助你掌握编译器优化技术。

目录
  1. 可视化演示
  2. torch.onnx.export导出流程与核心参数
  3. 动态维度与动态batch处理
  4. Opset版本选择与算子兼容性处理
  5. 不支持算子的自定义导出方法
  6. onnxsim模型简化与导出后验证
  7. AI 应用场景
  8. Python 演示
  9. 常见问答
  10. 常见误区
  11. 参考来源

可视化演示

展示图形渲染管线中光栅化、着色插值与深度测试的核心阶段，动态演示三角形从顶点到像素的完整渲染过程。

本节配有交互动效（进入视口后自动加载，约 3KB）。

三角形光栅化 顶点着色器输出三角形顶点坐标，光栅化阶段将连续几何转换为离散像素片段 扫描线算法逐行填充三角形内部，每个像素生成一个片段送入后续处理 fragment = rasterize(triangle, pixel_grid)

torch.onnx.export导出流程与核心参数

PyTorch模型导出为ONNX格式的核心函数是torch.onnx.export，它采用trace模式（基于TorchScript的trace机制）工作。在trace模式下，函数会使用一组示例输入（args参数）运行模型一次，记录执行过程中经过的所有算子及其调用顺序，最终将这些算子序列转换为ONNX计算图。这意味着导出的模型实际上记录的是示例输入所经过的那一条执行路径，而非模型代码的全部逻辑分支。

args参数是导出流程中最关键的输入，它必须提供与模型实际推理时形状和类型一致的示例张量。对于多输入模型，args应传入一个元组，每个元素对应模型forward方法的一个参数。此外，input_names和output_names参数用于为输入输出节点命名，这些名称在后续推理引擎加载模型时用于绑定数据，命名规范直接影响部署代码的可读性和正确性。

opset_version参数决定了导出ONNX模型所使用的算子集版本，不同opset版本支持的算子种类和语义有所不同。较高的opset版本通常包含更多算子和更精确的语义定义，但需要推理引擎（如ONNX Runtime、TensorRT）对相应版本提供支持。实践中通常在满足算子需求的前提下选择推理引擎已验证支持的opset版本，常见选择为opset 11至17之间。

其他重要参数包括：do_constant_folding（是否在导出时执行常量折叠优化，默认为True）、verbose（是否打印导出过程的详细信息用于调试）、dynamic_axes（定义哪些维度是动态的，后续章节详述）。导出完成后建议使用onnx.load加载模型并用onnx.checker.check_model验证模型结构的完整性。

动态维度与动态batch处理

在模型部署中，输入数据的batch size和序列长度等维度往往不是固定的，这就要求ONNX模型支持动态维度。torch.onnx.export通过dynamic_axes参数实现这一功能，该参数是一个字典，键为输入或输出张量的名称，值为一个字典，将维度索引映射到符号变量名称字符串。例如dynamic_axes={'input': {0: 'batch_size', 1: 'seq_len'}}表示input张量的第0维和第1维分别是可变的batch_size和seq_len。

符号变量名称（如'batch_size'、'seq_len'）本身只是标签，相同名称的符号变量在不同输入输出之间会建立关联约束。例如，如果input的第0维和output的第0维都映射到'batch_size'，则推理时这两个维度必须保持一致。这种机制保证了模型内部维度的一致性，同时允许跨不同推理请求灵活变化。

动态batch处理是最常见的动态维度应用场景。将batch维度设为动态后，同一个ONNX模型可以处理batch_size=1的在线推理请求，也可以处理batch_size=64的批量推理请求，无需为每种batch size导出单独的模型。这在服务端部署中极大简化了模型管理。

需要注意的是，并非所有算子都支持动态维度。某些算子在导出时可能会将维度信息硬编码到模型结构中（如view、reshape操作中通过计算得出的维度），导致虽然dynamic_axes声明了动态维度，但实际推理时改变该维度会报错。因此，导出后必须使用不同形状的输入进行多次推理测试来验证动态维度的正确性。

Opset版本选择与算子兼容性处理

ONNX的opset版本定义了当前ONNX规范支持的所有算子及其语义。每个PyTorch算子在导出时需要映射到特定opset版本中对应的ONNX算子，如果所选opset版本不包含该算子的映射定义，导出将失败。因此，opset版本的选择首先取决于模型中使用的算子类型——必须确保所有算子在目标opset版本中都有定义。

选择opset版本时，必须考虑目标推理引擎的支持范围。ONNX Runtime每个版本支持的opset范围不同，TensorRT对opset的支持也有版本限制。实践中的推荐策略是：先确认推理引擎支持的最高opset版本，然后在该版本下尝试导出，如果导出失败则逐步降低opset版本，直到找到所有算子都能正确映射的版本。

opset版本过低会导致某些新算子无法导出（如LayerNormalization在opset 17才原生支持），而opset版本过高可能导致推理引擎不支持或行为不一致。常见的兼容性问题是：在开发环境使用高版本opset成功导出，但部署环境的推理引擎版本较低无法加载。解决方案是锁定部署环境的推理引擎版本，并在该版本对应的opset范围内进行导出和测试。

onnx.checker.check_model是导出后的必要验证步骤，它会检查模型的计算图结构是否合法、张量类型是否一致、属性是否完整。此外，onnx.version_converter可以在不同opset版本之间转换模型，但并非所有算子都支持版本转换。建议在CI/CD流程中加入自动化的opset兼容性测试，确保模型在目标部署环境中的可用性。

不支持算子的自定义导出方法

当模型中包含PyTorch到ONNX没有标准映射的算子时，导出会失败。处理这种情况的首选策略是模型改造：用功能等价的ONNX标准算子组合替换不支持的算子。例如，某些自定义的归一化操作可以用标准的Mean、Sub、Pow、ReduceMean、Sqrt、Div等基本算子的组合来实现。这种方式虽然需要修改模型代码，但导出后的模型兼容性最好。

对于无法用标准算子替换的复杂操作，可以注册自定义的ONNX符号化函数（symbolic function）。通过torch.onnx.register_custom_op_symbolic或在新版API中使用torch.onnx.register_custom_export_symbolic，可以为自定义算子定义导出到ONNX的映射逻辑。自定义算子在ONNX模型中会以CustomOp的形式存在，推理时需要在推理引擎中注册对应的实现。

Script模式（torch.jit.script）是trace模式的补充方案。trace模式只能记录示例输入经过的执行路径，无法处理数据依赖的控制流（如if语句中条件依赖于输入张量的值）。script模式通过编译Python代码为TorchScript IR，能够正确导出包含控制流的模型。使用torch.jit.script(model)将模型编译后，再通过torch.onnx.export传入scripted model进行导出。

在实际工程中，混合使用trace和script是常见做法：对模型中不含控制流的部分使用trace导出，对包含复杂控制流的子模块使用script导出，然后通过torch.onnx.export的opset_version和dynamo选项灵活组合。PyTorch 2.0引入的torch.export和dynamo编译器进一步简化了这一流程，能够自动处理大部分控制流场景的导出问题。

onnxsim模型简化与导出后验证

PyTorch导出的ONNX模型通常包含大量冗余操作，如常量传播可以预计算的中间结果、重复的reshape/transpose操作、未使用的分支等。onnxsim（ONNX Simplifier）工具通过常量折叠（Constant Folding）、冗余算子消除（Redundant Operator Elimination）和代数化简（Algebraic Simplification）等优化手段，将模型简化为功能等价但结构更紧凑的形式。使用方式非常简单：onnxsim input_model.onnx output_model.onnx。

常量折叠是onnxsim最核心的优化之一。它将可以在编译时确定结果的算子序列（如形状已知的reshape、常量之间的算术运算）预先计算并替换为常量节点，显著减少推理时的计算量。冗余消除则合并功能重复的连续算子，例如连续的两次transpose如果互为逆操作则直接消除。

导出后的模型必须经过推理验证以确保与原始PyTorch模型的行为一致。标准做法是使用onnxruntime加载ONNX模型，用与导出时相同的输入数据进行推理，然后比较ONNX推理结果与PyTorch推理结果的数值差异。通常使用numpy.allclose或计算最大绝对误差（max absolute error）和相对误差（relative error）来衡量。对于浮点模型，误差在1e-5量级通常可接受；对于量化模型，误差容忍度需要适当放宽。

Netron是一款开源的神经网络模型可视化工具，支持ONNX、PyTorch、TensorFlow等多种格式。通过Netron可以直观地查看ONNX模型的完整计算图结构，检查每个节点的算子类型、输入输出张量的形状和名称、属性参数等。在调试导出问题时，Netron的可视化能力可以帮助快速定位异常节点，例如形状不匹配的reshape操作或名称错误的输入输出端口。Netron提供Web版（netron.app）和桌面客户端，是ONNX模型调试的必备工具。

AI 应用场景

BERT模型导出为ONNX用于推理加速

阶段：模型部署

解决的问题：将训练好的BERT文本分类模型导出为ONNX格式，使用ONNX Runtime进行推理加速，支持动态序列长度和动态batch size

具体案例：使用torch.onnx.export导出BERT模型，设置dynamic_axes={'input_ids':{0:'batch',1:'seq'}, 'attention_mask':{0:'batch',1:'seq'}, 'output':{0:'batch'}}，opset_version设为14，导出后用onnxsim简化，最后用onnxruntime.InferenceSession加载并进行性能benchmark

自定义算子模型导出到TensorRT部署

阶段：模型部署

解决的问题：模型包含自定义的ROI Align算子，标准ONNX导出失败，需要通过注册自定义symbolic函数或替换为标准算子组合来完成导出

具体案例：在目标检测模型中，自定义ROI Align算子通过torch.onnx.register_custom_op_symbolic注册导出映射，将其转换为ONNX的RoiAlign算子（opset 16），然后在TensorRT中加载时注册对应的plugin实现

模型导出后精度验证与问题排查

阶段：模型测试

解决的问题：导出ONNX模型后推理结果与PyTorch原始模型存在较大偏差，需要定位是导出参数问题、算子不支持还是浮点精度问题

具体案例：使用onnxruntime加载导出的模型，输入相同的测试数据，逐层对比中间输出与PyTorch模型的差异。发现某自定义LayerNorm实现因trace模式只走了固定epsilon分支导致精度问题，改用script模式重新导出解决

Python 演示

python59 行
import torchimport torch.onnximport onnximport onnxruntime as ortimport numpy as npfrom onnxsim import simplify # 1. 定义模型（示例：简单的分类模型）class SimpleClassifier(torch.nn.Module):    def __init__(self, input_dim, hidden_dim, num_classes):        super().__init__()        self.fc1 = torch.nn.Linear(input_dim, hidden_dim)        self.relu = torch.nn.ReLU()        self.fc2 = torch.nn.Linear(hidden_dim, num_classes)     def forward(self, x):        x = self.relu(self.fc1(x))        return self.fc2(x) # 2. 创建模型并设为评估模式model = SimpleClassifier(128, 64, 10)model.eval() # 3. 准备示例输入dummy_input = torch.randn(1, 128) # 4. 导出ONNX模型（含动态batch维度）torch.onnx.export(    model,    dummy_input,    "model.onnx",    opset_version=13,    input_names=["input"],    output_names=["output"],    dynamic_axes={        "input": {0: "batch_size"},        "output": {0: "batch_size"}    },    do_constant_folding=True) # 5. 验证ONNX模型结构onnx_model = onnx.load("model.onnx")onnx.checker.check_model(onnx_model) # 6. onnxsim简化onnx_model_sim, check = simplify(onnx_model)assert check, "onnxsim验证失败"onnx.save(onnx_model_sim, "model_sim.onnx") # 7. 使用ONNX Runtime推理验证sess = ort.InferenceSession("model_sim.onnx")test_input = np.random.randn(4, 128).astype(np.float32)onnx_output = sess.run(None, {"input": test_input})[0] # 8. 与PyTorch结果对比with torch.no_grad():    pt_output = model(torch.from_numpy(test_input)).numpy()print(f"最大绝对误差: {np.max(np.abs(pt_output - onnx_output)):.2e}")
目的：完整演示PyTorch模型导出为ONNX的全流程：模型定义、导出、动态维度设置、结构验证、onnxsim简化、ONNX Runtime推理验证及精度对比
  1. model.eval()导出前必须将模型设为评估模式，确保Dropout、BatchNorm等层的行为与推理时一致
  2. dummy_input = torch.randn(1, 128)示例输入的形状决定了导出模型中非动态维度的固定值，batch维度后续通过dynamic_axes设为动态
  3. dynamic_axes={"input": {0: "batch_size"}}将输入张量的第0维标记为动态batch_size维度，推理时可以传入任意batch大小的输入
  4. onnx.checker.check_model(onnx_model)验证导出的ONNX模型结构合法性，检查算子属性、张量形状推断等是否正确
  5. onnx_model_sim, check = simplify(onnx_model)使用onnxsim进行常量折叠和冗余消除，返回简化后的模型和验证结果
  6. np.max(np.abs(pt_output - onnx_output))计算PyTorch与ONNX Runtime推理结果的最大绝对误差，验证导出精度

常见问答

torch.onnx.export使用的是什么模式来记录模型计算图？
trace模式（基于TorchScript的trace机制），通过运行一次示例输入来记录执行路径上的所有算子。
dynamic_axes参数的作用是什么？
用于指定ONNX模型中哪些输入/输出张量的哪些维度是动态的（可变长度的），通过字典形式将维度索引映射到符号变量名称。
opset_version参数为什么重要？
它决定了导出模型使用的ONNX算子集版本，不同版本支持的算子种类和语义不同，直接影响模型的兼容性和功能完整性。
为什么导出前需要调用model.eval()？
因为eval模式会关闭Dropout、将BatchNorm切换为推理模式（使用running mean/var），确保导出模型的行为与推理时一致。
trace模式的主要局限是什么？
trace模式只能记录示例输入实际经过的执行路径，无法处理数据依赖的控制流（如if条件依赖于输入值），也无法记录循环等动态结构。
onnxsim的主要优化手段有哪些？
主要包括常量折叠（预计算常量表达式）、冗余算子消除（合并重复操作）、代数化简（简化等价表达式）等。
如何处理ONNX导出时遇到不支持的算子？
可以通过三种方式处理：1）用标准ONNX算子组合替换；2）注册自定义symbolic函数；3）使用script模式导出。
onnx.checker.check_model检查什么内容？
检查模型的计算图结构合法性，包括算子属性完整性、张量类型一致性、形状推断正确性、输入输出定义完整性等。
script模式和trace模式的区别是什么？
trace模式通过执行一次记录算子序列，无法处理数据依赖控制流；script模式通过编译Python代码为IR，能正确处理条件分支和循环等控制流结构。
导出ONNX后如何验证模型精度？
使用onnxruntime加载ONNX模型，输入与PyTorch相同的测试数据，比较两者推理输出的数值差异，通常用最大绝对误差或相对误差衡量。
input_names和output_names参数的作用是什么？
为ONNX模型的输入和输出节点命名，在推理引擎加载模型时通过这些名称绑定输入数据和获取输出结果。
do_constant_folding参数的作用是什么？
控制导出时是否执行常量折叠优化，设为True时会将可以预计算的常量表达式在导出阶段就计算好，减少推理时的计算量。
Netron工具在ONNX调试中有什么作用？
Netron可以可视化ONNX模型的完整计算图，帮助检查节点算子类型、张量形状、属性参数等，快速定位导出异常。
为什么dynamic_axes声明了动态维度但推理时仍可能报错？
因为某些算子（如view、reshape）可能在导出时将维度信息硬编码到模型结构中，导致虽然声明了动态维度但实际计算时无法适应变化的形状。
如何选择ONNX opset版本？
先确认目标推理引擎支持的最高opset版本，然后在该版本下尝试导出，如果失败则逐步降低版本直到所有算子都能正确映射。
ONNX模型中的符号变量名称有什么作用？
符号变量名称用于建立不同输入输出张量动态维度之间的关联约束，相同名称的符号变量在推理时必须保持一致的维度大小。

常见误区

误区导出模型与PyTorch行为完全一致

正解trace模式只记录示例输入经过的执行路径，如果模型包含数据依赖的控制流（如if/else分支），不同输入可能走不同路径，导出的模型只能复现示例输入走过的路径行为，其他路径的逻辑不会被导出。

为何常见：用户容易认为导出过程会完整转换模型的所有逻辑，忽略了trace模式的本质是'录制一次执行'而非'翻译全部代码'。

误区dynamic_axes后所有维度可任意变化

正解dynamic_axes只是告诉推理引擎哪些维度可以变化，但模型结构本身存在隐含约束。例如reshape操作中如果某个维度是通过其他维度计算得出的（如-1自动推断），则相关维度之间存在数学约束关系，不能独立任意变化。

为何常见：用户看到dynamic_axes允许维度变化，就认为可以随意设置任意值，忽略了模型内部算子之间的维度依赖关系和数学约束。

误区onnxsim简化不影响推理结果

正解onnxsim的优化（如常量折叠、算子融合、代数化简）在数学上等价，但由于浮点运算的不可结合性，运算顺序的改变可能导致微小的数值差异（通常在1e-6到1e-7量级）。对于精度敏感的场景，简化后仍需进行精度验证。

为何常见：用户认为'简化'只是结构优化不改变计算结果，忽略了浮点运算重排带来的精度影响，尤其在深层网络中微小差异可能累积。

参考来源