3、ONNX模型基础：理解ONNX格式、模型结构、输入输出张量

3.1 核心原理：ONNX 到底是什么

ONNX（Open Neural Network Exchange）是一个开放式的模型交换标准，它定义了一套统一的计算图格式，让不同深度学习框架（PyTorch、TensorFlow、PaddlePaddle 等）训练出的模型可以相互转换和部署。你可以把 ONNX 理解为“模型界的通用语言”——无论源框架是什么，只要导出为 ONNX 格式，就能被支持 ONNX 的推理引擎（如 ONNX Runtime、TensorRT、OpenVINO）加载和执行。

ONNX 模型本质上是一个有向无环图（DAG），由以下核心元素构成：

• 节点（Node）：图中的每个算子（如 Conv、Relu、MatMul），包含算子类型、属性（如卷积核大小、步长）和输入输出名称。
• 张量（Tensor）：数据以多维数组形式流动，每个张量有名称、数据类型（float32、int64 等）和形状（shape）。
• 图（Graph）：由节点和连接节点的张量边组成，包含模型的全部计算逻辑。
• 元数据（Metadata）：模型版本、生产者名称、Opset 版本等信息，用于兼容性管理。

下面这张图展示了 ONNX 模型文件内部的逻辑结构，以及输入张量如何经过算子节点最终产生输出张量：

从图中可以看到，输入张量 data_0 依次经过 Conv → Relu → MaxPool → Gemm 等节点，中间产生多个中间张量，最终得到输出张量 prob_1。权重参数作为节点的附加输入参与计算，而整个图的执行顺序由拓扑排序决定。

3.2 方法流程：如何用 Python 查看 ONNX 模型结构

在实际开发中，我们通常使用 onnx 库来加载和检查模型。核心流程如下：

1. 加载模型：使用 onnx.load() 读取 .onnx 文件。
2. 检查模型合法性：调用 onnx.checker.check_model() 验证模型是否符合规范。
3. 解析图结构：遍历 model.graph.node 获取算子列表。
4. 查看输入输出：读取 model.graph.input 和 model.graph.output 获取张量信息。

以下是一个完整的示例代码：

import onnx

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

# 2. 检查模型合法性
onnx.checker.check_model(model)

# 3. 获取图对象
graph = model.graph

# 4. 打印模型元信息
print(f"模型名称: {model.name}")
print(f"Opset 版本: {model.opset_import[0].version}")
print(f"生产者: {model.producer_name}")

# 5. 查看输入张量
print("\n=== 输入张量 ===")
for inp in graph.input:
    shape = [dim.dim_value for dim in inp.type.tensor_type.shape.dim]
    print(f"名称: {inp.name}, 形状: {shape}, 类型: {inp.type.tensor_type.elem_type}")

# 6. 查看输出张量
print("\n=== 输出张量 ===")
for out in graph.output:
    shape = [dim.dim_value for dim in out.type.tensor_type.shape.dim]
    print(f"名称: {out.name}, 形状: {shape}, 类型: {out.type.tensor_type.elem_type}")

# 7. 统计算子类型
from collections import Counter
op_types = Counter(node.op_type for node in graph.node)
print("\n=== 算子统计 ===")
for op, cnt in op_types.most_common(10):
    print(f"{op}: {cnt} 次")

运行上述代码后，你会看到类似下面的输出：

模型名称: resnet50
Opset 版本: 17
生产者: PyTorch

=== 输入张量 ===
名称: input, 形状: [1, 3, 224, 224], 类型: 1

=== 输出张量 ===
名称: output, 形状: [1, 1000], 类型: 1

=== 算子统计 ===
Conv: 53 次
Relu: 49 次
BatchNormalization: 53 次
MaxPool: 1 次
Gemm: 1 次
Flatten: 1 次

其中类型编号 1 对应 TensorProto.FLOAT，即 float32 类型。通过这种方式，你可以快速了解模型的输入输出规格和计算结构。

3.3 案例场景：从 PyTorch 导出并验证 ONNX 模型

假设你训练好了一个简单的 PyTorch 分类模型，现在需要将其部署到 ONNX Runtime。完整流程如下：

import torch
import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(2)
        self.fc = nn.Linear(16 * 16 * 16, 10)  # 假设输入 32×32

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = x.view(x.size(0), -1)
        return self.fc(x)

# 创建模型实例并导出
model = SimpleCNN()
model.eval()

dummy_input = torch.randn(1, 3, 32, 32)
torch.onnx.export(
    model,
    dummy_input,
    "simple_cnn.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)
print("模型已导出为 simple_cnn.onnx")

import onnxruntime as ort
import numpy as np

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

# 准备输入数据（模拟一张图片）
input_data = np.random.randn(1, 3, 32, 32).astype(np.float32)

# 执行推理
outputs = sess.run(None, {"input": input_data})
print(f"输出形状: {outputs[0].shape}")
print(f"预测类别: {np.argmax(outputs[0], axis=1)}")

这里的关键点在于：导出时的 dummy_input 形状必须与真实输入一致，否则导出的模型可能无法正确处理实际数据。使用 dynamic_axes 可以让 batch 维度可变，增强模型的灵活性。

3.4 误区与复盘：常见问题及解决思路

常见误区	问题表现	正确做法
忽略输入张量的数据布局	PyTorch 使用 NCHW，而某些框架使用 NHWC，导致推理结果错误	导出前确认输入布局，必要时使用 transpose 调整
动态维度处理不当	固定形状模型无法处理不同 batch 的输入	使用 dynamic_axes 指定可变维度
算子不支持	某些自定义算子无法导出，或目标推理引擎不支持	使用 torch.onnx.export 的 opset_version 参数，或注册自定义算子
精度下降	导出后模型精度与原始模型有差异	检查是否启用了 training=False，并对比中间层输出

理解 ONNX 的模型结构是后续进行模型优化、量化和部署的基础。掌握输入输出张量的查看方法，能帮助你快速定位问题，并为编写高效的推理代码打下坚实基础。