📚 课程目录第2章 ONNX基础

2、ONNX基础：ONNX格式介绍、模型导出与加载、运行时选择

2.1 ONNX格式的核心设计

ONNX（Open Neural Network Exchange）本质上是一种跨框架的模型中间表示协议。它不负责训练，也不直接执行推理，而是定义了一套与框架无关的计算图描述规范。一个完整的ONNX模型文件（通常以 .onnx 为后缀）内部采用Protobuf序列化存储，主要包含以下三个核心部分：

计算图（Graph）：描述数据从输入到输出的张量流动过程，由节点（Node）和边（Tensor）构成。每个节点代表一个算子（Op），例如 Conv、Relu、MatMul。
元数据（Metadata）：记录模型的版本、生产者（如PyTorch、TensorFlow）、算子集版本（Opset Version）等信息，用于兼容性判断。
权重参数（Initializer）：以张量形式内嵌在模型中的训练参数，如卷积核权重和偏置。

理解ONNX的关键在于：它只定义“图的结构”，而“如何执行”则交给不同的运行时后端。这意味着同一个 .onnx 文件，既可以被CPU上的 ONNX Runtime 加载，也可以被GPU上的TensorRT或OpenVINO优化执行。

核心要点：ONNX是“通用语言”，而非“执行引擎”。它解决了模型从训练框架到推理引擎的“最后一公里”格式转换问题，但本身不提供高性能计算能力。

2.2 模型导出：从训练框架到ONNX

以PyTorch为例，导出ONNX的核心API是 torch.onnx.export。导出过程并非简单的“格式转换”，而是通过追踪（Tracing）或脚本化（Scripting）方式记录模型的实际计算路径。最常用的是Tracing模式，它通过输入一个示例张量，让模型前向传播一次，从而记录下实际执行过的算子序列。

import torch
import torchvision.models as models

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

# 2. 构造示例输入（batch_size=1, 3通道, 224x224）
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出ONNX
torch.onnx.export(
    model,                      # 要导出的模型
    dummy_input,                # 示例输入
    "resnet18.onnx",            # 输出文件路径
    export_params=True,         # 导出权重参数
    opset_version=17,           # ONNX算子集版本
    do_constant_folding=True,   # 常量折叠优化
    input_names=["input"],      # 输入节点名称
    output_names=["output"],    # 输出节点名称
    dynamic_axes={              # 动态维度（可选）
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)
print("导出完成：resnet18.onnx")

导出后，建议立即使用 onnx.checker 和 onnx.shape_inference 验证模型完整性：

import onnx

# 加载并检查模型
model = onnx.load("resnet18.onnx")
onnx.checker.check_model(model)  # 抛出异常则说明模型结构有误

# 推断并打印模型信息
onnx.shape_inference.infer_shapes(model)
print(f"算子数量: {len(model.graph.node)}")
print(f"输入: {model.graph.input[0].name}, 形状: {model.graph.input[0].type.tensor_type.shape}")

导出建议：对于包含动态控制流（如循环、条件分支）的模型，Tracing可能无法完整捕获逻辑，此时应优先考虑 torch.jit.script 或使用 torch.onnx.export 的 scripting 方式。此外，opset_version 并非越高越好，需确认目标运行时支持的最高版本。

2.3 模型加载与运行时选择

加载ONNX模型并执行推理，最常用的运行时是 ONNX Runtime（简称ORT）。它由微软维护，支持CPU、CUDA、TensorRT、OpenVINO等多种执行提供程序（Execution Provider，简称EP）。选择不同EP的核心逻辑是：将ONNX图中的算子分配到可用的硬件加速器上执行。

运行时/EP	适用硬件	优势场景	备注
CPUExecutionProvider	任意x86/ARM CPU	通用兼容、无GPU环境	默认EP，支持INT8量化
CUDAExecutionProvider	NVIDIA GPU	标准CNN、Transformer	需安装CUDA/cuDNN
TensorrtExecutionProvider	NVIDIA GPU (TensorRT)	高吞吐、低延迟批量推理	需TensorRT库，首次构建引擎慢
OpenVINOExecutionProvider	Intel CPU/集成显卡	Intel平台优化、边缘部署	对Intel硬件友好

在Python中使用ORT加载模型并指定EP的典型流程如下：

import onnxruntime as ort
import numpy as np

# 1. 创建推理会话，按优先级指定EP列表
providers = [
    "TensorrtExecutionProvider",  # 优先使用TensorRT
    "CUDAExecutionProvider",      # 其次CUDA
    "CPUExecutionProvider"        # 最后CPU兜底
]
session = ort.InferenceSession("resnet18.onnx", providers=providers)

# 2. 查看实际启用的EP
print("当前使用的EP:", session.get_providers())

# 3. 构造输入并推理
input_name = session.get_inputs()[0].name
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
outputs = session.run(None, {input_name: input_data})
print("输出张量形状:", outputs[0].shape)

2.4 常见误区与复盘

误区一：盲目追求高opset版本。高版本opset可能包含新算子，但旧版ORT或TensorRT可能不支持，导致加载失败。建议先确认目标部署环境的ORT版本，再选择合适的opset。

误区二：忽略动态轴设置。如果模型需要处理不同batch size的输入，但导出时未设置 dynamic_axes，则运行时输入形状被固定，无法灵活调整。批量推理优化时，动态batch是常见需求。

复盘案例：某团队在将YOLOv5模型导出为ONNX后，直接使用默认CPU EP进行批量推理，发现吞吐量仅为GPU的1/20。排查后发现：

未安装CUDA EP，导致模型完全运行在CPU上；
模型输出包含大量后处理算子（如NMS），这些算子在GPU上效率反而低。

解决方案：将模型拆分为“主干特征提取”和“后处理”两部分，前者使用TensorRT EP加速，后者保留在CPU上执行，最终整体吞吐量提升了12倍。这个案例说明：运行时选择不是“全有或全无”，而是可以按算子粒度混合使用多个EP。