3、ONNX基础：ONNX是什么、为什么选择ONNX作为中间格式、ONNX模型的导出（PyTorch/TensorFlow）、ONNX模型可视化

各位同学，咱们今天聊聊ONNX。说实话，我刚开始接触模型部署那会儿，也被各种框架的模型格式搞得头大。PyTorch有.pt，TensorFlow有.pb，还有Keras的.h5……每个框架都像在说自己的方言。直到我遇到了ONNX，才算是找到了一个通用的“普通话”。

3.1 ONNX是什么？

ONNX，全称是Open Neural Network Exchange，直译过来就是“开放神经网络交换格式”。你把它想象成模型界的“通用语言”就行。

我个人的理解是：ONNX定义了一套标准的计算图表示方法。不管你是用PyTorch训练，还是用TensorFlow训练，最终都能导出成一个.onnx文件。这个文件里记录了网络的结构、每一层的参数、权重，以及数据流动的路径。

说白了，ONNX就是一个中间表示（Intermediate Representation，简称IR）。它不负责训练，只负责把训练好的模型“翻译”成一种标准格式，方便后续的推理引擎（比如咱们要讲的TensorRT）去读取和执行。

核心要点：ONNX = 模型结构的通用描述语言 + 权重的标准存储格式。

3.2 为什么选择ONNX作为中间格式？

你可能会问：我直接用PyTorch的模型不行吗？为什么非要转一手ONNX？

嗯，这里有几个非常现实的原因，我在项目里都踩过坑：

  • 框架解耦：训练框架和推理框架可以分开选。比如你用PyTorch训练，但推理时想用TensorRT加速。没有ONNX，你就得自己写转换代码，那工作量……想想都头疼。
  • 算子标准化：不同框架对同一个算子的实现可能不一样。ONNX定义了一套标准算子集，比如Conv、Relu、BatchNormalization等。导出时，框架会把自家的算子映射到ONNX标准算子上。这就保证了模型的可移植性。
  • 优化空间大：很多推理引擎（包括TensorRT）都对ONNX做了深度优化。它们会解析ONNX图，然后做图优化、算子融合、内存复用等操作。你直接喂一个原始框架的模型，优化效果往往不如先转ONNX。
  • 生态支持好：现在主流框架都支持导出ONNX，主流推理引擎也都支持导入ONNX。你想想看，这就像USB接口一样，成了事实上的标准。

我的经验：我曾经在一个项目里，客户要求用TensorRT部署，但他们的模型是用PaddlePaddle训练的。当时如果没有ONNX，我可能得手动重写整个网络结构。有了ONNX，一行命令就搞定了。所以，ONNX真的是部署工程师的“瑞士军刀”。

3.3 ONNX模型的导出

好，理论说完了，咱们来点实际的。怎么把训练好的模型导出成ONNX？

3.3.1 PyTorch模型导出ONNX

PyTorch官方提供了torch.onnx.export函数，用起来非常方便。我给你看个例子：

import torch
import torchvision.models as models

# 1. 加载预训练模型
model = models.resnet18(pretrained=True)
model.eval()  # 一定要切换到eval模式！

# 2. 创建一个 dummy input（虚拟输入）
#    注意：输入的形状要和模型要求的一致
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出ONNX
torch.onnx.export(
    model,               # 要导出的模型
    dummy_input,         # 模型的输入张量
    "resnet18.onnx",     # 输出的ONNX文件名
    export_params=True,  # 是否导出权重参数
    opset_version=11,    # ONNX算子集版本，我一般用11或13
    do_constant_folding=True,  # 是否做常量折叠优化
    input_names=['input'],     # 输入节点的名字
    output_names=['output'],   # 输出节点的名字
    dynamic_axes={
        'input': {0: 'batch_size'},   # 动态batch
        'output': {0: 'batch_size'}
    }
)

print("ONNX模型导出成功！")

避坑指南：我曾经因为忘记调用model.eval()，导致导出的ONNX模型里包含了BatchNormalization的running_mean和running_var的更新逻辑。结果推理时结果总是不对。所以，导出前一定要确保模型在eval模式。

3.3.2 TensorFlow模型导出ONNX

TensorFlow这边稍微麻烦一点，因为TensorFlow的模型格式有好几种（SavedModel、Keras H5、Frozen Graph等）。不过别担心，我们有tf2onnx这个工具。

先安装：pip install tf2onnx

然后，假设你有一个Keras模型：

import tensorflow as tf
import tf2onnx

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

# 2. 指定输入签名（input signature）
#    这一步是为了告诉tf2onnx输入的形状和数据类型
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)

# 3. 导出ONNX
model_proto, _ = tf2onnx.convert.from_keras(
    model, 
    input_signature=spec, 
    opset=13,
    output_path="my_model.onnx"
)

print("TensorFlow模型导出ONNX成功！")

如果你用的是TensorFlow的SavedModel格式，也可以用命令行工具：

python -m tf2onnx.convert \
    --saved-model ./saved_model_dir \
    --output model.onnx \
    --opset 13

我的建议：导出时，opset_version（或opset）不要选太新，也不要太旧。我一般选11到15之间。太新的算子集，有些推理引擎可能不支持；太旧的，又可能缺少一些高级算子。

3.4 ONNX模型可视化

模型导出后，你肯定想看看它长什么样。ONNX模型本质上是一个计算图，可视化能帮你快速检查网络结构是否正确。

我常用的工具有两个：

  • Netron：一个在线可视化工具，直接把.onnx文件拖进去就行。界面非常友好，支持查看每一层的输入输出形状、参数、甚至权重值。我个人强烈推荐。
  • ONNX GraphSurgeon：如果你需要编程式地查看或修改ONNX图，可以用这个库。它提供了Python API，可以遍历节点、修改属性、删除节点等。

用Netron看模型，就像看一张电路图。你能清晰地看到数据从输入层开始，经过卷积、池化、全连接，最后到输出层。如果发现某个节点的形状不对，或者多了奇怪的节点，那就要回去检查导出过程了。

小技巧：我每次导出ONNX后，第一件事就是用Netron打开看一眼。确认输入输出形状正确，确认没有多余的节点（比如一些训练时才需要的节点）。这一步能帮你省下后面排查问题的大量时间。

好了，关于ONNX的基础知识就讲到这里。总结一下：ONNX是模型部署的“通用语言”，它让你从框架的束缚中解放出来。导出时注意eval模式、opset版本、输入签名这些细节。导出后用Netron看一眼，心里就有底了。

下一章，咱们会深入讲解如何用TensorRT读取这个ONNX模型，并把它转换成真正的推理引擎。到时候你会发现，前面这些准备工作，都是值得的。

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