📚 课程目录第3章 ONNX模型基础

3、ONNX模型基础：理解ONNX格式、计算图、节点与张量，以及如何导出或获取ONNX模型

在RKNN模型转换的完整工具链中，ONNX（Open Neural Network Exchange）扮演着“中间表示层”的关键角色。无论是PyTorch、TensorFlow还是PaddlePaddle训练的模型，最终都需要统一转换为ONNX格式，才能被RKNN-Toolkit2解析并生成RKNN运行时模型。因此，深入理解ONNX的底层数据结构——计算图（Graph）、节点（Node）与张量（Tensor），是掌握模型转换、排查精度损失和算子兼容性问题的前提。本章将系统拆解ONNX的格式规范、核心概念，并给出从主流框架导出ONNX的完整操作路径与实战验证方法。

3.1 ONNX格式本质：一种可交互的中间表示协议

ONNX并非一种新的深度学习框架，而是一种开放、跨框架的模型表示规范。它使用Protobuf（Protocol Buffers）序列化格式存储模型数据，文件扩展名为.onnx。从RKNN转换的角度看，ONNX文件的价值在于它提供了一种“标准化”的模型描述语言，使得不同框架训练的模型能够被统一的解析器读取。一个完整的ONNX模型文件包含三个核心部分：模型元数据（ModelProto）、计算图（GraphProto）以及算子集（OperatorSet）。其中，计算图是模型的主体，它由节点（NodeProto）和张量（ValueInfoProto/TensorProto）构成，节点定义了计算操作，张量则承载了数据流。

核心原理：ONNX采用“有向无环图（DAG）”来描述神经网络。图中的每个节点代表一个算子（如Conv、Relu、Gemm），节点之间的有向边则通过张量名（Tensor Name）来隐式连接。这种设计使得ONNX模型具有以下特性：

框架无关性：模型不绑定任何特定训练框架的运行时。
可移植性：同一份ONNX文件可以在不同推理引擎（如ONNX Runtime、RKNN、TensorRT）中执行。
可读性：通过Netron等可视化工具，可以直观查看网络结构。

在RKNN转换流程中，ONNX文件是输入格式之一。RKNN-Toolkit2会解析ONNX计算图，将每个节点映射到NPU支持的算子，并优化张量的内存布局。因此，理解ONNX的图结构，实际上就是理解RKNN转换器的工作对象。

3.2 计算图、节点与张量的深度拆解

为了更精准地进行模型转换与排障，我们需要深入理解ONNX内部三个核心概念的具体表现形式。下面通过一个典型的卷积神经网络片段来剖析。

3.2.1 计算图（Graph）：模型的骨架

计算图是ONNX模型的主体容器。它定义了模型的输入、输出以及所有中间计算步骤。在GraphProto中，主要包含以下字段：

node：节点列表，按拓扑顺序排列。
input：模型的输入张量定义（名称、数据类型、维度）。
output：模型的输出张量定义。
initializer：初始器，存放模型的权重和偏置等常量张量。

计算图的执行逻辑是：从输入张量开始，依次执行每个节点，将节点的输出张量传递给下一个节点，直到产生最终输出。RKNN转换器会遍历这张图，识别每个节点的算子类型，并构建内部的算子映射表。

3.2.2 节点（Node）：计算单元

每个节点代表一个数学操作。节点定义包含三个关键属性：op_type（算子类型）、input（输入张量名列表）、output（输出张量名列表）。此外，节点还可以包含attribute（属性），用于传递算子参数，例如卷积的strides、pads、dilations等。以下是一个Conv节点的ONNX文本表示示例：

node {
  input: "input_1"        // 输入特征图
  input: "conv1_weight"   // 卷积核权重
  input: "conv1_bias"     // 偏置
  output: "conv1_output"  // 输出特征图
  op_type: "Conv"
  attribute {
    name: "strides"
    ints: 2
    ints: 2
  }
  attribute {
    name: "pads"
    ints: 1
    ints: 1
    ints: 1
    ints: 1
  }
}

在RKNN转换中，节点的op_type决定了转换器是否能够找到对应的NPU算子。如果遇到不支持的算子（如某些自定义算子），转换会报错或回退到CPU执行。

3.2.3 张量（Tensor）：数据载体

张量是ONNX中数据的基本单位。它可以是模型的输入、输出，也可以是节点之间的中间结果。张量定义包含name（名称）、elem_type（数据类型）、shape（维度信息）。对于权重等常量数据，则存储在initializer中，以TensorProto格式保存原始字节数据。张量的数据布局（如NCHW或NHWC）对RKNN转换至关重要，因为NPU对数据布局有特定要求，转换器通常需要插入Transpose节点来调整布局。

概念	ONNX中的表示	对RKNN转换的影响
计算图	GraphProto，包含节点列表与张量定义	决定转换器遍历顺序与优化策略
节点	NodeProto，op_type + 输入输出张量名	算子映射的关键依据，不支持的算子导致转换失败
张量	ValueInfoProto / TensorProto	数据布局、精度（FP32/FP16/INT8）影响NPU执行效率
初始器	存储权重、偏置等常量	权重会被转换为RKNN内部格式，需注意内存对齐

3.3 从主流框架导出ONNX模型：操作流程与实战

获取ONNX模型通常有两种途径：一是从公开模型库（如ONNX Model Zoo）直接下载，二是从自己训练的框架中导出。在实际RKNN开发中，第二种途径更为常见。下面分别介绍PyTorch和TensorFlow的导出方法。

3.3.1 PyTorch导出ONNX（推荐路径）

PyTorch通过torch.onnx.export函数实现动态图到静态ONNX图的转换。导出时，需要提供一个示例输入，用于追踪模型的计算图。关键参数包括opset_version（算子集版本）、input_names和output_names。以下是一个标准导出流程：

import torch
import torchvision.models as models

# 1. 加载模型并设置为评估模式
model = models.resnet18(pretrained=True)
model.eval()

# 2. 构造示例输入（与训练时数据布局一致）
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出ONNX
torch.onnx.export(
    model,                      # 要导出的模型
    dummy_input,                # 示例输入
    "resnet18.onnx",            # 输出文件路径
    export_params=True,         # 导出权重
    opset_version=12,           # 推荐使用12或更高版本
    do_constant_folding=True,   # 常量折叠优化
    input_names=["input"],      # 输入张量名
    output_names=["output"],    # 输出张量名
    dynamic_axes={              # 动态维度（可选）
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)
print("ONNX模型导出成功")

经验提醒：在导出时，务必设置opset_version为RKNN-Toolkit2支持的版本（通常为12~15）。过高的opset版本可能导致RKNN解析器不兼容。另外，如果模型包含nn.Upsample等算子，建议使用opset_version=12以上，以获得更稳定的Resize算子表示。

3.3.2 TensorFlow导出ONNX（通过tf2onnx）

TensorFlow模型需要借助tf2onnx工具进行转换。首先将Keras模型或SavedModel保存为.pb文件，然后使用命令行工具转换。以下是一个Keras模型转换示例：

# 1. 保存Keras模型为SavedModel格式
model.save("saved_model/")

# 2. 使用tf2onnx命令行转换
python -m tf2onnx.convert \
    --saved-model saved_model/ \
    --output model.onnx \
    --opset 13 \
    --inputs-as-nchw input_1:0

# 3. 验证转换结果
python -c "import onnx; m=onnx.load('model.onnx'); onnx.checker.check_model(m); print('模型有效')"

注意，TensorFlow默认使用NHWC数据布局，而ONNX标准通常使用NCHW。在转换时，--inputs-as-nchw参数可以指定输入张量按NCHW处理，避免后续RKNN转换时出现布局混乱。

3.4 案例拆解：一个YOLOv5模型的ONNX导出与验证

为了将上述理论应用于实际，我们以YOLOv5s目标检测模型为例，完整演示从PyTorch导出ONNX并进行初步验证的过程。该案例在RKNN转换中非常典型，因为YOLOv5包含多输出、后处理算子等复杂结构。

背景与目标：我们需要将训练好的YOLOv5s模型部署到RK3588平台的NPU上。第一步是获得一个结构清晰、算子兼容的ONNX文件。

初始现象：直接使用YOLOv5官方仓库的export.py导出时，生成的ONNX文件包含大量自定义的Detect层算子，这些算子在RKNN中无法直接支持。

分析判断过程：通过Netron可视化ONNX图，发现Detect层被表示为一个大的Concat节点加上多个Sigmoid和Add操作。RKNN-Toolkit2虽然支持这些基础算子，但输出张量的组织方式与NPU后处理要求不一致。

实际操作步骤：

修改YOLOv5的export.py，设置--include onnx，并添加--opset 12参数。
导出后，使用onnx-simplifier工具简化模型，消除冗余的Shape和Gather操作。
手动修改ONNX图，将Detect层的输出拆分为三个独立的输出张量（分别对应80x80、40x40、20x20的特征图）。
使用onnx.checker验证模型完整性，并打印输入输出张量的形状。

结果对比：简化后的ONNX模型大小从14.8MB降至14.2MB，节点数量减少约15%。在后续RKNN转换中，算子映射成功率从82%提升至97%，仅剩少量Resize算子需要特殊处理。

风险与代价：手动修改ONNX图需要谨慎处理张量维度，一旦出错会导致输出形状不匹配。建议在修改后使用ONNX Runtime进行推理验证，确保输出与原始PyTorch模型一致。

风险提示：不要直接使用未经简化的ONNX模型进行RKNN转换。复杂的图结构（如大量Shape、Gather、Unsqueeze操作）会显著增加转换时间，并可能导致NPU内存分配失败。建议在导出后使用onnx-simplifier进行图优化。

3.5 常见误区与总结复盘

在ONNX模型准备阶段，开发者常陷入以下误区，导致后续RKNN转换受阻：

误区一：忽略opset版本兼容性。使用过高的opset版本（如15以上）导出的模型，可能包含RKNN-Toolkit2尚未支持的算子变体。建议固定使用opset 12或13。
误区二：不检查动态维度。如果导出的ONNX包含动态轴（如动态batch），RKNN转换时可能无法确定最优的内存布局。对于NPU部署，建议固定输入分辨率。
误区三：忽视数据布局。PyTorch默认NCHW，TensorFlow默认NHWC。在导出时未统一布局，会导致RKNN转换器插入大量不必要的Transpose节点，增加推理延迟。
误区四：跳过模型验证。导出后未使用ONNX Runtime进行数值比对，直接进行RKNN转换，导致精度问题难以定位。

总结复盘：ONNX模型是RKNN转换的“源头”。一个高质量的ONNX模型应具备以下特征：算子类型标准、图结构简洁、数据布局明确、输入输出定义清晰。在导出阶段投入时间进行优化，往往能避免后续RKNN转换中数倍的调试成本。建议在每次导出后，使用Netron可视化检查图结构，并使用onnx.checker和onnxruntime进行双重验证，确保模型在进入RKNN工具链之前是“干净”的。