📚 课程目录第2章 ONNX基础

ONNX基础：格式概述、模型结构、张量与算子、模型导入导出

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

ONNX（Open Neural Network Exchange）是一个开放的、跨框架的深度学习模型交换标准。它的核心价值在于定义了一套与框架无关的计算图描述格式，使得模型可以在 PyTorch、TensorFlow、PaddlePaddle 等不同训练框架之间迁移，并最终部署到各类推理引擎（如 ONNX Runtime、TensorRT、OpenVINO）中。你可以将 ONNX 理解为模型的“通用语言”——训练框架负责“写作”，推理引擎负责“阅读”，而 ONNX 文件就是那份双方都能理解的“手稿”。

从本质上看，ONNX 文件描述的是一个静态数据流图：图中的节点是算子（Operator），边是张量（Tensor）数据。它不包含训练所需的梯度信息或优化器状态，只保留前向推理所需的计算逻辑。这种精简设计使得 ONNX 文件体积小、结构清晰，非常适合作为端侧部署的中间表示。

关键认知：ONNX 不是推理引擎，也不是训练框架，而是一个中间表示层（IR）。它的职责是“准确描述计算”，而“如何高效执行”则交给下游的推理引擎去优化。

模型结构：计算图的三要素

一个标准的 ONNX 模型由以下三部分构成，它们共同定义了一张完整的计算图：

图（Graph）：包含所有节点、张量、输入输出和初始化数据，是模型的主体。
节点（Node）：每个节点代表一次算子调用，包含算子类型（op_type）、输入张量名列表、输出张量名列表以及算子属性（attribute）。
张量（Tensor）：图中的数据载体，包括模型输入、输出以及中间计算结果。权重和偏置等常量以初始化器（Initializer）的形式存储在图中。

此外，模型还包含元数据（Metadata），如 ir_version（IR版本）、producer_name（导出工具名称）和 opset_import（算子集版本），这些信息决定了模型使用的语法规范和算子版本兼容性。

张量与算子：模型的基本语言

张量（Tensor）是 ONNX 中所有数据的统一表示形式。一个张量由以下属性描述：

数据类型（data_type）：如 float32、int64、uint8 等。端侧部署中，float16 和 int8 量化类型尤为常见。
形状（shape）：如 [1, 3, 224, 224] 表示批量大小为 1、3 通道、224×224 像素的图像。ONNX 支持动态维度（用 None 或 symbol 表示），但端侧推理通常建议固定形状以换取性能优化。
数据存储：原始数据以线性数组形式存储，按行主序（row-major）排列。

算子（Operator）定义了张量之间的计算规则。每个算子有固定的 op_type（如 Conv、MatMul、Reshape），并通过属性（Attribute）来传递非张量参数（如卷积的 kernel_shape、strides、pads）。算子对张量的操作遵循严格的类型和形状推导规则，这保证了模型在导出时就能发现大部分结构错误。

算子类别	典型算子	端侧部署注意事项
张量变换	Reshape, Transpose, Concat, Split	Transpose 可能导致内存拷贝，尽量在导出前优化布局
卷积/池化	Conv, MaxPool, AveragePool	分组卷积、空洞卷积需确认推理引擎支持
激活函数	ReLU, Sigmoid, HardSwish	某些激活在 INT8 下需特殊实现，注意量化支持
矩阵运算	MatMul, Gemm, Einsum	Einsum 在部分端侧引擎上不支持，建议展开为 MatMul
归一化	BatchNormalization, LayerNormalization	BN 在推理时可折叠到 Conv 中，减少计算量
逻辑/比较	Greater, Where, NonZero	动态控制流算子会阻碍图优化，尽量避免

方法流程：模型的导入与导出

将训练好的模型转换为 ONNX 格式，通常遵循以下流程。以 PyTorch 为例：

准备模型：将模型切换到推理模式（model.eval()），并构造一个符合输入形状的示例张量。
执行导出：调用 torch.onnx.export，指定模型、示例输入、输出路径和算子集版本。
验证模型：使用 onnx.checker.check_model 检查模型结构合法性，再用 onnxruntime 跑一遍推理，对比输出与原始框架是否一致。
简化与优化：使用 onnx-simplifier 消除冗余节点，或手动调整图结构以适配目标端侧引擎。

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",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},  # 动态batch
    opset_version=17
)

# 4. 验证模型
import onnx
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX 模型导出并验证成功！")

对于 TensorFlow 模型，可使用 tf2onnx 工具；PaddlePaddle 则通过 paddle2onnx 完成转换。无论哪种框架，导出后都应立即进行数值一致性验证，确保输出误差在可接受范围内（通常 1e-4 以内）。

案例场景：MobileNetV3 的端侧部署准备

假设我们要将 MobileNetV3-Small 部署到手机端。在导出 ONNX 时，我们做了以下关键决策：

固定输入形状：将输入固定为 [1, 3, 224, 224]，避免动态维度带来的额外开销。
算子集版本选择：使用 opset 12，因为该版本对 HardSwish 和 HardSigmoid 的支持更成熟，且兼容大多数端侧引擎。
融合 BatchNormalization：在导出前将 BN 层与 Conv 层融合，减少模型中的节点数量。

导出后，我们使用 onnxruntime 进行推理测试，发现输出与 PyTorch 原始模型的最大误差为 2.3e-6，完全满足部署要求。随后，该 ONNX 文件被输入到 TensorRT 中进行 INT8 量化，最终在 Jetson Nano 上达到了 12ms 的推理延迟。

实操建议：在导出前，先检查模型中是否有自定义算子（如自定义损失函数中的操作）。如果存在，需要将其替换为 ONNX 标准算子，或注册自定义算子实现，否则导出会失败。

误区与复盘：常见问题及规避策略

在 ONNX 模型转换过程中，开发者常遇到以下问题，提前了解可有效避免踩坑：

常见误区	后果	正确做法
忽略 `model.eval()` 直接导出	模型中残留 Dropout、BatchNorm 的训练行为，推理结果错误	导出前务必切换到 eval 模式
使用动态轴但未指定 `dynamic_axes`	导出的模型仅支持固定 batch，部署时灵活性受限	明确声明需要动态的维度
算子集版本设置过高	端侧推理引擎不支持新算子，加载失败	根据目标引擎支持的 opset 版本选择
导出后不做数值验证	模型精度下降但无法定位是转换问题还是量化问题	用 ONNX Runtime 跑测试数据，对比输出差异
忽略输入数据的预处理对齐	端侧推理结果与训练时不一致（如归一化方式不同）	将预处理（mean/std 归一化）一并写入 ONNX 图或严格对齐

重要提醒：ONNX 模型只是“可运行”的中间表示，并不代表“高效运行”。从 ONNX 到端侧推理引擎，还需要经过图优化、算子融合、内存规划、量化等一系列步骤。不要期望导出的 ONNX 文件直接达到最优性能，它只是整个优化链路的第一步。

复盘一个典型失败案例：某团队导出的 YOLOv5 模型在 PC 端 ONNX Runtime 上运行正常，但部署到手机端时崩溃。排查发现，模型中使用了 torch.where 导出的 Where 算子，而该手机端引擎的旧版本不支持此算子。解决方案是将 Where 替换为 Mul + Add 的组合，重新导出后问题解决。这提醒我们：在导出前，必须确认目标推理引擎的算子支持列表，必要时对模型结构进行针对性调整。