2. ONNX模型基础

核心原理：ONNX 是什么，为什么需要它

ONNX（Open Neural Network Exchange）是一个开放的、跨框架的模型交换标准。它的核心价值在于：让一个在 PyTorch 中训练好的模型，能够被导出为 ONNX 格式，再被 TensorFlow、PaddlePaddle、OpenVINO、ONNX Runtime 等其他框架或推理引擎加载和执行。这种“一次导出，多处部署”的能力，解决了深度学习模型在生产环境中依赖特定训练框架的痛点。

从本质上看，ONNX 定义了一套与框架无关的计算图描述规范。它不关心模型是如何训练出来的，只关心模型在推理时的数据流和计算逻辑。因此，ONNX 文件可以被视为一个“中间表示（IR）”，它位于训练框架与推理引擎之间，起到翻译官和桥梁的作用。

ONNX 的典型工作流：

  1. 在 PyTorch / TensorFlow 中完成模型训练。
  2. 调用框架的导出 API（如 torch.onnx.export）将模型转换为 .onnx 文件。
  3. 使用 ONNX Runtime 或其他推理引擎加载该文件，进行推理或进一步优化。
PyTorch 训练模型 TensorFlow 导出 ONNX 模型 .onnx 中间表示 加载 ONNX Runtime CPU / CUDA / TensorRT OpenVINO Intel 硬件加速 其他支持 ONNX 的推理引擎

模型结构：计算图、节点与张量

ONNX 模型的核心是一个有向无环图（DAG），它描述了数据从输入到输出的流动过程。理解这个图结构，是后续进行算子融合、精度调优和性能分析的基础。

一个 ONNX 文件主要由以下三部分构成：

  • 节点（Node）：图中的每个圆圈代表一个算子（Op），例如 Conv、Relu、MatMul、Add。节点定义了计算操作的类型和属性（如卷积的 kernel_shape、strides）。
  • 张量（Tensor）：节点之间的有向边传递的是多维数组，即张量。张量有名称、数据类型（如 float32、int64）和形状（shape）。在 ONNX 中，张量通常以 ValueInfo 的形式在图中声明。
  • 图属性（Graph）：包含图的输入、输出、节点列表以及初始化的权重张量（initializer）。权重通常作为图的输入之一，但在推理时被固定。

关键概念：ONNX 中的 initializer 是模型的权重参数。它们被存储在图的 initializer 列表中，并在推理时作为常量输入传递给节点。在模型优化时，我们经常会将某些节点（如 BatchNormalization）的权重“折叠”到相邻的卷积节点中，以减少计算量。

方法流程：使用 Netron 可视化模型结构

面对一个陌生的 ONNX 模型，最有效的入门方式是使用可视化工具查看其结构。Netron 是目前最流行的 ONNX 可视化工具，它支持 Windows、macOS、Linux 以及浏览器在线版。

使用 Netron 的基本流程：

  1. 打开 Netron（本地应用或 https://netron.app）。
  2. 将 .onnx 文件拖拽到窗口中。
  3. 点击任意节点，右侧面板会显示该节点的属性（如卷积核大小）、输入张量的名称与形状、输出张量的形状。
  4. 点击输入或输出张量，可以查看其数据类型和维度信息。

除了 Netron，命令行工具 onnx.shape_inference 也可以帮助我们快速检查模型的张量形状是否正确，这在排查导出错误时非常有用。

# 使用 Python API 查看模型基本信息
import onnx

model = onnx.load("model.onnx")
graph = model.graph

print("模型输入:")
for inp in graph.input:
    print(f"  - {inp.name}, 形状: {[dim.dim_value for dim in inp.type.tensor_type.shape.dim]}")

print("模型输出:")
for out in graph.output:
    print(f"  - {out.name}, 形状: {[dim.dim_value for dim in out.type.tensor_type.shape.dim]}")

print(f"节点数量: {len(graph.node)}")
print(f"初始权重数量: {len(graph.initializer)}")

案例/场景：一个 MiniCNN 模型的 ONNX 结构剖析

为了更直观地理解上述概念，我们以一个包含卷积、批归一化、ReLU 和全连接层的小型 CNN 模型为例。假设我们已将其导出为 minicnn.onnx，使用 Netron 打开后，可以看到如下结构：

节点序号 算子类型 输入张量 输出张量 关键属性
1 Conv input (1,3,32,32), conv1.weight conv1_output (1,16,32,32) kernel_shape: [3,3], strides: [1,1]
2 BatchNormalization conv1_output, bn1.weight, bn1.bias bn1_output (1,16,32,32) epsilon: 1e-5
3 Relu bn1_output relu1_output (1,16,32,32) 无
4 MaxPool relu1_output pool1_output (1,16,16,16) kernel_shape: [2,2], strides: [2,2]
5 Flatten pool1_output flatten_output (1, 4096) 无
6 Gemm flatten_output, fc1.weight, fc1.bias output (1, 10) transB: 1

通过这张表，我们可以清晰地看到数据从 (1,3,32,32) 的输入图像，逐步经过卷积、归一化、激活、池化，最终变为 (1,10) 的分类向量的过程。在调优时，如果我们发现 Conv 算子的耗时占比过高，就可以针对性地考虑是否将其与 BatchNormalization 融合，或者尝试不同的卷积算法。

常见误区：初学者容易混淆 ONNX 节点与 PyTorch 层的概念。PyTorch 中的 nn.Sequential 可能包含多个操作，但导出为 ONNX 后，这些操作会被拆分成更细粒度的节点。例如，一个 nn.Linear 层在 ONNX 中可能表现为 Gemm 节点，也可能被拆分为 MatMul 和 Add 两个节点，这取决于导出时的算子集版本和设置。

误区/复盘：模型可视化时的常见陷阱

在查看 ONNX 模型结构时，有几个容易踩的坑，需要特别注意：

  • 忽略输入输出的显式声明：ONNX 图的输入不仅包含原始数据（如 image），还包含所有未初始化的权重。如果某个权重没有被放入 initializer，它就会出现在 graph.input 中，导致推理时需要额外传入参数。使用 onnx.checker.check_model 可以验证模型是否规范。
  • 动态轴与静态轴的混淆：在导出时，如果 batch size 被固定为 1，那么输入张量的形状就是 (1,3,32,32)。若希望支持动态 batch，需要在导出时指定动态轴。在 Netron 中，动态轴通常显示为 None 或 symbol，而不是具体的数字。
  • 只关注节点类型，忽略张量形状：调优时，不仅要看有哪些算子，更要关注张量在节点之间的形状变化。例如，一个 Reshape 或 Transpose 操作可能会导致张量在内存中不连续，从而引发额外的拷贝开销。在 Netron 中点击张量边，可以查看其每一步的形状变化。

复盘一个实际案例：某次模型转换后推理速度反而变慢，通过 Netron 检查发现，模型中出现了大量 Cast 节点（将 float32 转为 float64），原因是 PyTorch 模型中某些层的权重被意外地转换为了 double 类型。将这些权重修正为 float32 后，速度恢复正常。这提醒我们，可视化不仅是看结构，更是排查性能问题的第一把钥匙。

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