3、ONNX基础：ONNX标准介绍、为什么选择ONNX作为中间表示、ONNX模型的导出与可视化

3.1 ONNX标准介绍

ONNX，全称Open Neural Network Exchange。说白了，它就是一套模型描述语言。

我刚开始接触深度学习框架时，最头疼的就是模型迁移。PyTorch训练好的模型，想部署到TensorRT上？得先转成Caffe，再转成TensorRT。中间各种算子不兼容，改来改去，一个模型能折腾好几天。

ONNX的出现，就是为了解决这个痛点。它定义了一套统一的算子集和计算图格式。不管你是用PyTorch、TensorFlow还是MXNet，只要导出成ONNX格式，就能在支持ONNX的推理引擎上跑起来。

ONNX的核心组件其实就两个：

  • 算子集（Operator Set）：定义了各种神经网络层对应的标准操作，比如Conv、Relu、BatchNormalization等。每个算子都有明确的输入输出规范和属性定义。
  • 计算图（Graph）：用节点（Node）表示算子，用边（Edge）表示张量流动。整个模型就是一个有向无环图。

嗯，这里要注意。ONNX并不是一个推理框架。它不负责执行模型，只负责描述模型。你可以把它理解成一种「模型通用语言」。

ONNX的版本演进：目前主流的是opset 11到opset 18。我个人建议，导出时尽量用opset 11或13，兼容性最好。opset 18虽然支持更多新算子，但很多推理引擎还没跟上。

3.2 为什么选择ONNX作为中间表示

你可能要问：市面上那么多中间格式，为什么偏偏选ONNX？

我在项目中试过直接转TensorRT，也试过用其他中间格式。最后发现，ONNX有几个不可替代的优势。

第一，生态最广。几乎所有主流框架都支持导出ONNX。PyTorch自带torch.onnx.export，TensorFlow有tf2onnx工具，连老牌的Caffe2都支持。你想想看，一个格式能打通这么多框架，省了多少事。

第二，工具链成熟。ONNX有配套的优化工具，比如onnx-simplifier可以简化计算图，onnxruntime可以直接做推理验证。我经常用onnxruntime跑一遍导出的模型，确认精度没问题，再交给TensorRT。

第三，调试方便。ONNX模型是结构化的，你可以用Netron工具直接可视化。哪层算子不对，哪条边断了，一眼就能看出来。我曾经排查过一个模型精度问题，就是通过可视化发现BatchNormalization的epsilon参数被错误地设成了0。

第四，标准化程度高。ONNX由微软、Facebook、亚马逊等公司共同维护。它的算子定义非常严谨，每个算子的行为都有明确规范。这保证了不同框架导出的ONNX模型，行为是一致的。

我的经验：如果你要部署到TensorRT，ONNX几乎是必经之路。TensorRT对ONNX的支持最好，官方提供了trtexec工具和onnx-tensorrt解析器。虽然也可以直接转Caffe或UFF，但坑多，不推荐。

3.3 ONNX模型的导出

导出ONNX模型，不同框架的步骤略有不同。我以PyTorch为例，这是我最常用的方式。

先看一个最简单的例子：

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 16, 3)
        self.relu = nn.ReLU()
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(16, 10)
    
    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x)
        x = self.pool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

model = SimpleModel()
model.eval()

# 构造一个虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出ONNX
torch.onnx.export(
    model,
    dummy_input,
    "simple_model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    },
    opset_version=13
)

这里有几个关键参数，我解释一下：

  • input_names / output_names：给输入输出张量起个名字。TensorRT解析时会用到这些名字。
  • dynamic_axes：指定哪些维度是动态的。比如batch_size，你训练时是32，部署时可能变成1或64。不设置的话，导出的模型就固定了batch_size。
  • opset_version：算子集版本。我一般用13，兼容性和功能都比较均衡。

我曾经踩过的坑：导出时忘记调用model.eval()。训练模式下，BatchNormalization和Dropout的行为跟推理时不一样。导出的ONNX模型精度会出问题。嗯，这个坑我至少踩过两次。

导出完成后，你可以用onnxruntime验证一下：

import onnxruntime as ort
import numpy as np

# 创建推理会话
sess = ort.InferenceSession("simple_model.onnx")

# 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 运行推理
outputs = sess.run(None, {"input": input_data})
print(outputs[0].shape)  # 应该输出 (1, 10)

如果这一步能跑通，说明ONNX模型导出成功了。接下来就可以交给TensorRT了。

3.4 ONNX模型的可视化

可视化ONNX模型，我推荐用Netron。这是一个开源的模型可视化工具，支持ONNX、TensorFlow、Caffe等多种格式。

使用方式很简单：

  • 在线版：打开 https://netron.app ，直接把.onnx文件拖进去。
  • 离线版：pip install netron，然后在代码里调用netron.start("model.onnx")，会自动在浏览器打开。

可视化之后，你能看到什么？

首先是整个计算图的拓扑结构。每个节点代表一个算子，每条边代表张量流动。你可以点击某个节点，查看它的详细属性，比如卷积的kernel_size、stride、padding等。

其次是输入输出的形状信息。Netron会显示每个张量的维度。如果某个张量的形状是动态的，它会显示为"dynamic"或具体的维度名称。

我通常用可视化来做两件事：

  1. 检查模型结构是否正确。比如有没有多余的节点，算子顺序对不对。
  2. 定位精度问题。如果TensorRT推理结果不对，我会对比ONNX模型和原始PyTorch模型的中间层输出。Netron能帮你快速找到哪一层开始出现偏差。

一个小技巧：如果导出的ONNX模型太大，可视化时很卡。你可以用onnx-simplifier先简化一下：pip install onnx-simplifier，然后运行python -m onnxsim input.onnx output.onnx。简化后的模型去掉了冗余的Shape、Gather等算子，看起来清爽多了。

另外，你也可以用Python代码来解析ONNX模型：

import onnx

# 加载模型
model = onnx.load("simple_model.onnx")

# 打印计算图信息
graph = model.graph
print(f"模型名称: {graph.name}")
print(f"输入节点数: {len(graph.input)}")
print(f"输出节点数: {len(graph.output)}")
print(f"算子节点数: {len(graph.node)}")

# 遍历所有算子
for node in graph.node:
    print(f"算子: {node.op_type}, 名称: {node.name}")
    print(f"  输入: {node.input}")
    print(f"  输出: {node.output}")

这种方式适合做自动化检查。比如你可以写个脚本，批量检查一批ONNX模型的结构是否合规。

好了，ONNX的基础知识就讲到这里。下一章我们会深入ONNX算子集，看看哪些算子TensorRT支持，哪些不支持，以及遇到不支持的算子该怎么办。

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