4、ONNX中间表示：ONNX是什么、为什么选择ONNX、PyTorch导出ONNX的基本流程

好，咱们进入第四章。这一章聊的是ONNX——一个在模型部署里绕不开的中间格式。

说实话，我刚开始做部署那会儿，对ONNX的理解就是“一个转换工具”。后来踩的坑多了，才慢慢明白它到底扮演了什么角色。今天我把这些经验掰开揉碎了讲给你听。

ONNX是什么？

ONNX，全称Open Neural Network Exchange。直译过来就是“开放神经网络交换格式”。

说白了，它就是一个模型界的“通用语言”。

你想想看，PyTorch训练出来的模型是.pt或.pth文件，TensorFlow的是.pb或.h5，Keras又是另一种。这些格式彼此不互通。你总不能为了部署一个PyTorch模型，专门去学一套TensorFlow的部署工具链吧？

ONNX就是来解决这个问题的。它定义了一套标准化的算子集合和计算图描述方式。不管你是用PyTorch、TensorFlow还是其他框架训练的模型，只要导出成ONNX格式，就能被各种推理引擎（比如TensorRT、ONNX Runtime、OpenVINO）读取和加速。

核心要点：ONNX不是推理引擎，也不是训练框架。它只是一个中间表示（Intermediate Representation, IR）。它的使命就是“桥接”。

我在项目中遇到过这样一个场景：客户要求用TensorRT部署，但团队里有人用PyTorch，有人用TensorFlow。最后我们统一导出ONNX，再转TensorRT。嗯，省了不少扯皮的功夫。

为什么选择ONNX？

你可能要问：为什么非得用ONNX？直接PyTorch转TensorRT不行吗？

行，但不够好。我给你列几个理由。

1. 框架解耦

ONNX让你从特定框架中解放出来。你训练时用PyTorch，部署时用TensorRT。中间这一层ONNX，让你不用关心TensorRT到底支不支持PyTorch的某个骚操作。

2. 算子标准化

每个框架都有自己的算子实现。PyTorch的Conv2d和TensorFlow的Conv2D，底层实现细节不同。ONNX定义了一套标准算子，比如Conv、Relu、Gemm等。导出时，框架会把自家算子映射到ONNX标准算子。这样下游工具只要支持ONNX算子集就行。

3. 生态成熟

现在主流推理引擎基本都支持ONNX。TensorRT、ONNX Runtime、OpenVINO、Triton Inference Server……你随便挑。而且ONNX社区活跃，新算子、新优化策略更新很快。

4. 可视化与调试

ONNX模型可以用Netron直接打开看结构。我经常用这个工具检查导出后的计算图有没有问题。比如某个算子被拆成了多个小算子，或者形状不对。一眼就能看出来。

个人建议：刚开始接触ONNX时，多看看Netron里的图。它能帮你建立“计算图”的直觉。我早期很多bug都是靠看图发现的。

PyTorch导出ONNX的基本流程

好，理论说完了，咱们上手实操。PyTorch导出ONNX，核心就一个函数：torch.onnx.export()。

流程其实很简单，但细节决定成败。我带你走一遍。

第一步：准备模型和输入

你得有一个训练好的模型，以及一个“样本输入”。这个样本输入不一定要是真实数据，但形状和数据类型必须和实际推理时一致。

import torch
import torchvision.models as models

# 加载一个预训练的ResNet18
model = models.resnet18(pretrained=True)
model.eval()  # 别忘了切换到eval模式

# 创建一个样本输入，batch_size=1，3通道，224x224
dummy_input = torch.randn(1, 3, 224, 224)

注意：一定要调用model.eval()。否则BatchNorm和Dropout的行为会不一样，导出的ONNX在推理时可能出问题。我曾经因为这个坑，排查了整整一个下午。

第二步：调用export函数

核心参数就几个：模型、样本输入、输出文件名、输入输出名称、动态轴（可选）。

torch.onnx.export(
    model,               # 要导出的模型
    dummy_input,         # 样本输入
    "resnet18.onnx",     # 输出文件名
    input_names=["input"],   # 输入节点名称
    output_names=["output"], # 输出节点名称
    dynamic_axes={
        "input": {0: "batch_size"},   # 第0维是动态的
        "output": {0: "batch_size"}
    },
    opset_version=17     # ONNX算子集版本
)

这里我重点说几个参数：

  • input_names / output_names：给输入输出张量起个名字。TensorRT里会用到这些名字来绑定数据。
  • dynamic_axes：指定哪些维度是动态的。比如batch_size不确定，就把它设成动态。注意，TensorRT对动态形状的支持有限，能固定尽量固定。
  • opset_version：ONNX算子集的版本。版本越高，支持的算子越多，但兼容性可能下降。我一般用当前PyTorch支持的默认版本，或者稍微新一点的。

第三步：验证导出的ONNX模型

导出完别急着走。一定要验证一下，确保模型结构和数值都对。

import onnx
import onnxruntime as ort

# 检查模型结构是否合法
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX模型结构检查通过！")

# 用ONNX Runtime跑一下推理，对比结果
ort_session = ort.InferenceSession("resnet18.onnx")
ort_inputs = {ort_session.get_inputs()[0].name: dummy_input.numpy()}
ort_outputs = ort_session.run(None, ort_inputs)

# 和PyTorch结果对比
with torch.no_grad():
    torch_output = model(dummy_input)

# 检查数值误差
diff = torch.abs(torch_output - torch.tensor(ort_outputs[0])).max()
print(f"最大数值误差：{diff.item():.6f}")

如果误差在1e-5以内，基本没问题。如果误差很大，可能是某些算子精度问题，或者模型里用了不支持的算子。

避坑指南：我曾经导出一个带自定义算子的模型，ONNX checker通过了，但推理结果完全不对。后来发现是自定义算子的梯度计算和ONNX映射有问题。所以，数值验证这步千万别省。

第四步：处理常见问题

导出过程中，你可能会遇到这些问题：

问题 原因 解决方案
不支持的操作 PyTorch的某些操作没有对应的ONNX算子 用torch.onnx.export的custom_opsets参数注册自定义算子，或者改写模型结构
动态形状失败 某些算子不支持动态形状 尽量固定形状，或者用torch.onnx.export的dynamic_axes精确指定
数值精度问题 FP16 vs FP32，或者算子实现差异 检查opset_version，或者用onnx-simplifier简化模型

小结

ONNX是模型部署的“通用语言”。它让你从框架绑定中解脱出来，让PyTorch模型能顺利进入TensorRT的优化流水线。

导出流程就四步：准备模型和输入 → 调用export → 验证结构 → 处理问题。每一步都有坑，但只要你按流程走，大部分问题都能提前发现。

下一章，咱们会深入ONNX的算子集和计算图结构。到时候你会更清楚，为什么有些模型导出来就是跑不动——嗯，那都是血泪教训换来的经验。

← 上一章 📖 返回目录 下一章 →