ONNX 模型导出

ONNX 是开放神经网络交换格式：把训练好的模型序列化为标准化计算图，让「训练在 PyTorch、推理在各后端」的跨框架链路成立。本篇覆盖 opset 概念、torch.onnx.export 导出流程、onnxruntime 推理验证与典型部署链路。

ONNX 定位

ONNX（Open Neural Network Exchange）用 protobuf 定义模型中间表示：一张有向无环计算图，节点是算子（Conv、MatMul、Relu……），边是张量，权重作为 initializer 内嵌于文件。任何框架实现「导出器」或「执行器」即可接入同一交换格式。

解决的问题：各框架原生格式互不兼容（见深度学习框架对比），而生产端推理后端众多——ONNX Runtime、TensorRT、OpenVINO、浏览器 WASM。训练侧模型权重保存的 .pt/.pth 是 PyTorch 私有 pickle 序列化，脱离 PyTorch 生态即不可用；ONNX 把「模型在哪训」与「模型在哪跑」解耦。

与推理模型格式的分工：safetensors/GGUF 面向 LLM 生态的权重分发与量化推理，仍依赖加载侧还原计算图；ONNX 是通用算子图标准，自带结构、覆盖 CV/语音/小模型全谱，两者交集在 transformers 导出（Optimum）。

opset：算子集版本

opset（operator set）是 ONNX 算子集版本号，决定图内可用算子及其语义。导出时指定 opset_version，模型会携带该版本；推理后端只支持到某个最高 opset。

opset 越新功能越全（如 opset 17 引入 LayerNormalization），但旧后端可能拒绝加载
导出器把 PyTorch 算子分解/映射为 ONNX 算子组合，映射缺失时报错或提示调整 opset

选型原则：按部署后端的支持上限定 opset，不盲目追新。

torch.onnx.export 导出

经典导出器基于 TorchScript trace：给模型一个 dummy input 跑一遍前向，记录实际执行的算子序列，固化成静态图。

import torch
 
model = MyModel()  # nn.Module，构建见 [[torch.nn]]
model.load_state_dict(torch.load("model.pth"))
model.eval()  # 必须切 eval：冻结 dropout 随机性与 BN 的 running stats
 
dummy = torch.randn(4, 3, 224, 224)  # 数值无意义，仅驱动一次前向 trace
 
torch.onnx.export(
    model, dummy, "model.onnx",
    input_names=["input"], output_names=["output"],
    opset_version=17,
    dynamic_axes={  # 声明符号维度，否则 batch 被固化为 4
        "input": {0: "batch"},
        "output": {0: "batch"},
    },
)

要点：

dummy input 的 shape/dtype/device 须与真实推理一致，具体数值任意
dynamic_axes 把 batch（乃至序列长度）标记为动态维度，避免每个 batch size 导出一份模型
trace 只记录走过的分支——数据依赖的控制流被静默丢弃，这正是 torch.autograd 动态图与部署静态图的根本冲突

新版后端 torch.onnx.export(..., dynamo=True) 走 TorchDynamo 路线：在 Python bytecode 层捕获 FX graph，追踪到 ATen 算子再 lowering 成 ONNX，对控制流、复杂数据结构支持更好；动态维度改用 dynamic_shapes 参数（torch.export 语法）。旧 TorchScript 导出器已宣布弃用，进入维护周期。

onnxruntime 验证

导出不等于正确，标准流程是结构校验加数值对齐：

import numpy as np
import onnx
import onnxruntime as ort
import torch
 
onnx.checker.check_model("model.onnx")  # 校验图连通性、算子签名、版本合法性
 
opts = ort.SessionOptions()
opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
sess = ort.InferenceSession("model.onnx", opts, providers=["CPUExecutionProvider"])
 
x = np.random.randn(8, 3, 224, 224).astype(np.float32)  # batch=8，检验动态轴生效
ort_out = sess.run(None, {"input": x})[0]
 
with torch.no_grad():
    pt_out = model(torch.from_numpy(x)).numpy()
 
np.testing.assert_allclose(ort_out, pt_out, rtol=1e-3, atol=1e-5)

容差来源：不同后端算子实现（卷积分块、gemm 累加顺序）的浮点误差，1e-3 量级相对误差属正常；差异巨大通常意味着 trace 走错分支或 dtype 隐式转换。SessionOptions 与 providers（CUDA/TensorRT EP）是生产侧性能 tuning 的入口。

部署链路与 TensorRT

PyTorch 训练 → .pth 权重 → torch.onnx.export → model.onnx
    ├─ ONNX Runtime 直接推理（CPU / GPU / 移动端）
    └─ TensorRT 解析 ONNX → 编译 engine → 极致 GPU 推理

TensorRT 不直接读 PyTorch 模型，ONNX 是其官方入口：trtexec --onnx=model.onnx 完成算子融合（Conv+BN+ReLU）、精度校准（FP16/INT8）与 kernel 自动调优，产出绑定特定 GPU 架构的 engine。代价是 engine 不可移植，换卡需重新编译。

局限

动态控制流：数据依赖的 if 与变长循环难以静态化，需展开或映射到 ONNX Loop 算子，支持有限
自定义算子：私有 CUDA kernel 无对应 ONNX 算子，须分解实现或注册 custom op（推理后端还要配套支持）
单向快照：导出即冻结的静态图，不能续训、不含优化器状态与训练逻辑；迭代回到 PyTorch 侧，改完重导
数值非严格等价：算子分解与替换在边缘情形引入微小差异

参考：ONNX 官网

知识花园

探索