ONNX基础：ONNX格式概述、计算图结构、节点与张量、模型文件组成

在深入PyTorch导出ONNX并处理动态shape之前，我们必须先建立对ONNX本身的扎实理解。ONNX（Open Neural Network Exchange）是一个开放的、跨框架的深度学习模型表示标准，它不隶属于任何单一框架，而是作为模型交换的“通用语言”。本章将系统拆解ONNX的格式本质、计算图结构、核心组成元素以及模型文件的物理组织方式，为后续动态shape的实战操作打下坚实的认知基础。

2.1 ONNX格式概述：为什么需要一种“中间表示”

在深度学习工程化实践中，我们经常面临一个现实问题：PyTorch训练好的模型，如何高效地部署到TensorRT、OpenVINO、ONNX Runtime或自研推理引擎中？直接使用PyTorch的TorchScript导出，往往受限于Python运行时和框架特定的算子实现。ONNX的核心价值在于它定义了一套框架无关的、可扩展的计算图描述规范。它将模型表示为一个有向无环图（DAG），图中的节点是算子（Op），边是张量（Tensor），从而将模型从特定框架的API中解放出来。

ONNX格式基于Protocol Buffers（Protobuf）序列化，这是一种高效的二进制编码机制，比JSON等文本格式更紧凑，解析速度更快。一个ONNX文件（通常以.onnx为后缀）本质上是一个ModelProto对象，它包含了模型的所有信息。理解这一点，是后续使用onnx库进行节点修改、shape校验和动态轴设定的前提。

2.2 计算图结构：从数据流视角理解模型

ONNX的计算图（Graph）是模型的核心逻辑载体。它由三个基本元素构成：图输入（Graph Input）、图输出（Graph Output）和节点（Node）。图输入定义了模型接收的外部数据入口，图输出定义了模型的最终结果。节点则代表具体的数学运算（如卷积、矩阵乘法、激活函数）。

数据在图中沿着边（即张量）流动。每个节点的输入张量要么来自图输入，要么来自前一个节点的输出。这种数据流模型使得ONNX天然支持静态图分析，例如shape推断、内存规划以及算子融合。对于动态shape处理而言，我们关注的是图输入和中间张量的dim_param字段，它允许将某个维度声明为符号变量（如"batch"），而不是固定的整数。

2.3 节点与张量：算子的具体化与数据载体

节点（Node）是计算图的基本执行单元。每个节点必须包含op_type（算子类型，如"Conv"、"Relu"）、input（输入张量名称列表）和output（输出张量名称列表）。节点还可以包含attribute，用于传递算子特定的参数，例如卷积的strides、pads，或者Reshape的目标shape。在动态shape场景中，某些attribute（如Reshape的shape）可能需要根据输入动态计算，此时该attribute通常会被转换为一个额外的输入张量。

张量（Tensor）是数据的具体载体。在ONNX中，张量不仅包括模型的权重参数（存放在GraphProto.initializer中），也包括数据流中的中间特征图。每个张量由ValueInfoProto描述，其中最关键的是TensorShapeProto。该结构包含一个dim列表，每个dim要么是固定的dim_value（整数），要么是符号化的dim_param（字符串）。例如，一个动态batch的输入可以表示为dim_param: "batch"，而通道数固定为3，则表示为dim_value: 3。

// 一个典型的动态shape输入描述（伪代码）
input {
  name: "input"
  type {
    tensor_type {
      elem_type: FLOAT
      shape {
        dim { dim_param: "batch" }   // 动态维度
        dim { dim_value: 3 }
        dim { dim_value: 224 }
        dim { dim_value: 224 }
      }
    }
  }
}

2.4 模型文件组成：从ModelProto到磁盘文件

一个完整的ONNX模型文件在逻辑上是一个ModelProto对象，其内部结构可以拆解为以下几个部分：

1. 元数据（Metadata）：包括ir_version（IR版本）、producer_name（导出框架）、opset_import（算子集版本）。这些信息决定了模型使用的语法规则和算子特性。
2. 计算图（Graph）：这是模型的主体，包含所有节点、初始器（权重）、图输入输出。
3. 初始器（Initializer）：直接嵌入在GraphProto中的权重张量列表。它们不通过节点计算产生，而是直接作为某些节点的输入。
4. 值信息（ValueInfo）：包括input、output和value_info（中间张量），用于声明张量的名称、数据类型和shape。

在磁盘上，这些结构被Protobuf序列化为二进制流。我们可以使用onnx.load()加载模型，然后通过Python API访问model.graph、model.opset_import等属性。理解文件组成有助于我们精准定位需要修改的字段，例如，当我们需要将某个输入的维度改为动态时，实际上是在修改graph.input中对应ValueInfoProto的shape.dim。

2.5 案例拆解：手动修改一个静态模型的batch维度为动态

为了将上述概念串联起来，我们通过一个具体案例来演示如何操作ONNX模型文件。假设我们有一个PyTorch导出的静态batch（batch=1）的模型，现在需要将其改为动态batch。

背景与目标：模型输入shape为[1, 3, 224, 224]，我们需要将其修改为["batch", 3, 224, 224]，以便在推理时接受任意batch大小的输入。

操作步骤：

import onnx

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

# 2. 获取图输入
graph_input = model.graph.input[0]

# 3. 修改第一个维度为动态
dim = graph_input.type.tensor_type.shape.dim[0]
dim.dim_param = "batch"   # 将dim_value替换为dim_param
# 注意：需要清除原有的dim_value
del dim.dim_value

# 4. 保存修改后的模型
onnx.save(model, "dynamic_model.onnx")

# 5. 验证模型
onnx.checker.check_model(model)
print("模型修改成功，输入shape:", graph_input.type.tensor_type.shape)

分析判断过程：在修改过程中，我们只修改了graph.input中的第一个dim。但模型内部可能存在多个节点依赖输入shape，例如Reshape或Flatten操作。如果这些节点使用了固定的1作为batch维度，模型在推理时就会出错。因此，修改后必须使用onnxruntime进行实际推理测试，或者使用onnx-simplifier等工具进行优化。

结果对比与风险：修改后，模型可以接受[8, 3, 224, 224]的输入。但风险在于，如果模型内部有Flatten操作将[1, 3, 224, 224]展平为[1, 150528]，那么当batch变为8时，展平后的shape应为[8, 150528]。ONNX Runtime会自动处理这种基于符号维度的传播，但前提是算子支持动态shape。部分算子（如某些自定义Op）可能不支持，需要额外处理。

2.6 总结复盘：ONNX基础要点回顾

本章我们系统梳理了ONNX的核心概念。从格式概述中，我们理解了ONNX作为中间表示的价值；从计算图结构中，我们掌握了图、节点、张量三者之间的关系；从节点与张量的细节中，我们明确了动态shape在底层是如何通过dim_param实现的；从模型文件组成中，我们知道了修改模型本质上是修改ModelProto的特定字段。

在后续章节中，我们将深入PyTorch的torch.onnx.export函数，学习如何通过dynamic_axes参数在导出阶段就生成动态shape模型。但无论工具如何高级，理解本章的底层数据结构都是排查问题的基础。例如，当导出后的模型shape不符合预期时，你需要能够打开.onnx文件，检查graph.input和graph.output的dim_param是否正确，以及中间节点是否引入了不必要的固定shape约束。