4. PyTorch 到 ONNX 导出实战
在 PyTorch/ONNX 到 MLIR 的转换链路中，torch.onnx.export 是承上启下的关键一步。它把 PyTorch 的动态图执行结果固化为 ONNX 的静态计算图，后续的 ONNX 优化、MLIR 导入与 lowering 都建立在这份导出产物之上。导出质量直接决定后续转换能否成功，因此本节围绕参数配置、动态轴、自定义算子和验证四个维度展开。
一、核心原理：导出做了什么
torch.onnx.export 的本质是“追踪 + 图构建”。它使用示例输入执行模型，通过 TorchScript 追踪记录算子调用序列，再将其翻译为 ONNX 节点，并写入 graph、initializer、value_info 等结构。因此，示例输入的形状、类型、设备会直接影响导出图的输入签名与部分常量折叠结果。
导出过程可概括为：准备模型与示例输入 → 设置 opset 与动态轴 → 执行追踪 → 算子映射与图构建 → 序列化写入 .onnx 文件。理解这条链路，才能明白为什么“同样的模型，换一个示例输入，导出结果可能不同”。
二、参数详解：最常调整的几项
导出参数众多，但真正影响后续 MLIR 转换的主要集中在以下几项：
| 参数 | 作用 | 实践建议 |
|---|---|---|
opset_version | 指定 ONNX 算子集版本 | 与推理后端、MLIR 导入器对齐，常用 13～17 |
input_names / output_names | 命名输入输出张量 | 必须显式命名，便于后续定位与调试 |
dynamic_axes | 声明动态维度 | 批大小、序列长度等变化维度必须声明 |
do_constant_folding | 是否常量折叠 | 默认开启，可减少图规模，但会隐藏部分结构 |
training | 导出训练/推理模式 | 转换场景固定为 TrainingMode.EVAL |
export_params | 是否导出权重 | 默认 True，权重写入 initializer |
verbose | 打印导出细节 | 排障时开启，观察算子映射过程 |
torch.onnx.export(
model,
dummy_input,
"model.onnx",
opset_version=17,
input_names=["input_ids", "attention_mask"],
output_names=["logits"],
dynamic_axes={
"input_ids": {0: "batch", 1: "seq_len"},
"attention_mask": {0: "batch", 1: "seq_len"},
"logits": {0: "batch", 1: "seq_len"},
},
do_constant_folding=True,
training=torch.onnx.TrainingMode.EVAL,
)三、动态轴设置：最容易踩坑的地方
动态轴通过 dynamic_axes 声明，格式为 {张量名: {维度索引: 维度名}}。维度名只是符号标记，真正起作用的是“该维度不再被固定为示例输入中的具体数值”。
常见误区有三类。第一，只声明输入不声明输出，导致输出仍被固定为示例形状，后续推理遇到不同 batch 直接报错。第二，维度索引写错，例如把 seq_len 标到了 batch 维上，图能导出但语义错误。第三，对本身必须静态的维度（如类别数）也声明为动态，反而让 MLIR 侧无法做形状推导和内存规划。
四、自定义算子处理
当模型包含 PyTorch 自定义算子或 ONNX 尚未覆盖的算子时，导出通常会失败或产生 UnsupportedOperatorError。处理思路有三条：
- 注册符号映射：通过
torch.onnx.register_custom_op_symbolic把自定义算子映射到已有 ONNX 算子组合。 - 导出为自定义域算子：使用
custom_opsets声明自定义域，让算子以domain::Op形式保留在图中，后续在 MLIR 侧实现 lowering。 - 图重写替换：在导出前用
torch.fx或模块替换，把自定义算子改写为等价的标准算子序列。
from torch.onnx import register_custom_op_symbolic
def my_gelu(g, input, approximate):
return g.op("Gelu", input, approximate_s=approximate)
register_custom_op_symbolic("my_ops::gelu", my_gelu, opset_version=17)五、导出验证：三道检查
导出完成不等于正确。建议按以下顺序验证：
- 结构检查：
onnx.checker.check_model确认图结构合法、opset 与节点属性匹配。 - 形状检查：用
onnx.shape_inference.infer_shapes推断并打印输入输出形状，确认动态轴标记生效。 - 数值检查：用 ONNX Runtime 加载模型，与 PyTorch 输出逐元素比对，误差应控制在
1e-4量级以内。
import onnx, onnxruntime, numpy as np
onnx.checker.check_model("model.onnx")
model = onnx.load("model.onnx")
onnx.shape_inference.infer_shapes(model)
sess = onnxruntime.InferenceSession("model.onnx")
onnx_out = sess.run(None, {"input_ids": x.numpy()})[0]
torch_out = model(x).detach().numpy()
print("max diff:", np.abs(onnx_out - torch_out).max())六、复盘：导出阶段的典型问题
实践中，导出失败大多集中在四类：动态轴声明不完整导致形状固化；opset 版本过低导致算子无法映射；自定义算子未注册符号；以及 do_constant_folding 把带控制流的子图折叠成常量，掩盖了原本的结构。排障时优先开启 verbose=True，观察具体是哪个算子映射失败，再针对性处理。
把导出这一步做扎实，后续 ONNX 到 MLIR 的导入才能拿到干净、形状语义明确、算子可追溯的计算图，这也是整条转换链路稳定性的第一道保障。