一、ONNX 基础：为什么它成了模型转换的“通用语言”

做模型部署这些年，我接触过各种各样的框架。PyTorch 训练好的模型，想搬到 TensorRT 上跑？或者从 TensorFlow 切到 ONNX Runtime？说白了，最头疼的就是格式不统一。ONNX 的出现，很大程度上解决了这个痛点。

我个人习惯把 ONNX 看作模型界的“普通话”。不管你是用 PyTorch 还是 TensorFlow，只要导出成 ONNX，就能在支持 ONNX 的推理引擎上跑。嗯，这背后其实有一套标准化的设计思路。

1.1 ONNX 格式到底长什么样？

ONNX 全称是 Open Neural Network Exchange，微软和 Meta 联合搞的。它定义了一套统一的计算图表示。你想想看，一个神经网络无非就是一堆算子（Conv、ReLU、Add 这些）连来连去。ONNX 就是把这些算子和它们之间的数据流动，用 protobuf 序列化成一个文件。

一个 ONNX 模型文件，核心包含三部分：

  • 模型结构（GraphProto）：定义了计算图，包括所有节点（算子）、输入输出张量的名称和形状。
  • 权重参数（TensorProto）：所有训练好的权重，以二进制形式存着。
  • 元数据（ModelProto）：比如算子集版本、模型版本、生产者信息等。

举个例子，一个简单的 Conv+ReLU 网络，ONNX 内部大概是这样描述的：

ir_version: 8
graph {
  node {
    input: "input"
    input: "conv_weight"
    input: "conv_bias"
    output: "conv_output"
    op_type: "Conv"
  }
  node {
    input: "conv_output"
    output: "relu_output"
    op_type: "Relu"
  }
  input {
    name: "input"
    type { tensor_type { elem_type: 1 shape { dim { dim_value: 1 } dim { dim_value: 3 } dim { dim_value: 224 } dim { dim_value: 224 } } } }
  }
  output {
    name: "relu_output"
    type { tensor_type { elem_type: 1 shape { dim { dim_value: 1 } dim { dim_value: 64 } dim { dim_value: 224 } dim { dim_value: 224 } } } }
  }
  initializer {
    name: "conv_weight"
    data_type: 1
    dims: 64
    dims: 3
    dims: 3
    dims: 3
    raw_data: "..."
  }
}

你看，每个节点就是一个算子，输入输出都是张量。权重存在 initializer 里。这种结构清晰明了，解析起来也快。

1.2 为什么选择 ONNX 作为中间表示？

我在项目中遇到过好几次这样的场景：团队用 PyTorch 做实验，但生产环境是 TensorRT。如果没有 ONNX，就得手写模型转换脚本，或者用框架自带的转换工具，那叫一个痛苦。

ONNX 作为中间表示，有几个实实在在的好处：

  • 框架解耦：训练和推理可以分开选型。你爱用 PyTorch 训练就用 PyTorch，推理端想用 TensorRT 还是 ONNX Runtime 都行。
  • 算子标准化：ONNX 定义了一套标准算子集（比如 opset 版本）。不同框架的算子，只要映射到 ONNX 算子，就能互通。
  • 工具链成熟：ONNX 有配套的优化工具（onnx-simplifier、onnxoptimizer），还有调试工具（Netron 可视化）。
  • 硬件适配好：NVIDIA、Intel、AMD 这些硬件厂商，都优先支持 ONNX 导入。说白了，你用 ONNX，硬件兼容性基本不用愁。

核心观点：ONNX 不是最快的推理引擎，但它是最好的“桥梁”。模型从训练到部署，ONNX 是最省心的中间格式。

1.3 常用模型导出为 ONNX

下面我分别说说 PyTorch 和 TensorFlow 怎么导出 ONNX。这两个框架的导出方式差别挺大，但核心思路是一样的：把动态图或静态图，转成 ONNX 的静态计算图。

1.3.1 PyTorch 导出 ONNX

PyTorch 用的是 torch.onnx.export。这里有个关键点：PyTorch 是动态图，导出时需要给一个 dummy input，让模型跑一遍，trace 出计算图。

import torch
import torchvision.models as models

# 加载一个预训练模型
model = models.resnet18(pretrained=True)
model.eval()

# 构造一个 dummy input，形状要和模型输入一致
dummy_input = torch.randn(1, 3, 224, 224)

# 导出 ONNX
torch.onnx.export(
    model,               # 模型
    dummy_input,         # 示例输入
    "resnet18.onnx",     # 输出文件名
    export_params=True,  # 是否导出权重
    opset_version=17,    # ONNX 算子集版本
    do_constant_folding=True,  # 常量折叠优化
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)

print("导出成功！")

这里有个坑，我曾经踩过：dynamic_axes 参数。如果你不指定，导出的 ONNX 模型输入形状是固定的（比如 batch size 只能是 1）。但实际部署时，你可能想支持可变 batch。所以，我建议你一开始就加上 dynamic_axes，哪怕你暂时只用固定 batch。

小技巧：导出后，用 Netron 打开 onnx 文件看一眼。如果图结构和你预期的不一样，八成是某个算子没被正确 trace。常见问题包括 if 分支、循环、或者用了 torch.no_grad 但没处理好。

1.3.2 TensorFlow 导出 ONNX

TensorFlow 导出 ONNX 稍微麻烦一点。因为 TensorFlow 有多个模型保存格式（SavedModel、H5、Checkpoint），而且 TF 的算子集和 ONNX 的映射关系没那么直接。

我一般用 tf2onnx 这个工具。它支持从 SavedModel 或 H5 文件直接转。

# 安装 tf2onnx
# pip install tf2onnx

# 命令行转换（从 SavedModel）
python -m tf2onnx.convert \
    --saved-model ./saved_model_dir \
    --output model.onnx \
    --opset 17

# 或者从 Keras H5 文件
python -m tf2onnx.convert \
    --keras ./model.h5 \
    --output model.onnx \
    --opset 17

如果你在代码里操作，也可以这样：

import tensorflow as tf
import tf2onnx

# 加载 Keras 模型
model = tf.keras.models.load_model("model.h5")

# 指定输入签名
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)

# 转换
output_path = "./model.onnx"
model_proto, _ = tf2onnx.convert.from_keras(model, input_signature=spec, opset=17)

# 保存
with open(output_path, "wb") as f:
    f.write(model_proto.SerializeToString())

避坑指南：我曾经遇到过一个情况，TF 模型里用了自定义层（比如自己写的 Lambda 层），tf2onnx 直接报错。解决办法是先把自定义层用标准 TF 算子重写，或者用 tf.function 包装一下。总之，ONNX 只认识标准算子，自定义算子需要提前注册或替换。

1.4 导出后的验证与调试

导出 ONNX 不是终点，验证才是。我每次导出后，都会做两件事：

  1. 用 onnxruntime 跑一遍推理，对比原始框架的输出。如果精度对不上，说明导出有问题。
  2. 用 onnx.checker 检查模型合法性，看看有没有结构错误。
import onnx
import onnxruntime as ort
import numpy as np

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

# 用 ONNX Runtime 推理
ort_session = ort.InferenceSession("resnet18.onnx")
input_name = ort_session.get_inputs()[0].name
output_name = ort_session.get_outputs()[0].name

# 构造输入
dummy_input = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 推理
outputs = ort_session.run([output_name], {input_name: dummy_input})
print("ONNX 推理输出形状:", outputs[0].shape)

嗯，这一步千万别省。我见过太多人导出 ONNX 后直接扔给 TensorRT，结果精度不对，排查半天才发现是导出时某个算子没映射好。

1.5 小结

ONNX 作为中间表示，最大的价值就是“一次导出，到处运行”。PyTorch 和 TensorFlow 的导出方式虽然不同，但核心思路一致：把模型计算图标准化。我个人建议，无论你用什么框架训练，都养成导出 ONNX 的习惯。这不仅是部署的需要，也是模型格式备份的好方式。

下一章，我会深入讲 ONNX 的算子集和 opset 版本管理。这东西看着枯燥，但实际部署时经常遇到版本不兼容的坑。到时候我会拿几个真实案例给你拆解。

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