2. ONNX模型基础

在深入GPU推理加速之前，必须先理解ONNX模型本身。ONNX（Open Neural Network Exchange）是一个开放的模型表示标准，它不绑定任何特定框架，而是定义了一套统一的计算图描述方式。你可以把ONNX模型看作一个“中间语言”——PyTorch、TensorFlow等框架训练出的模型，都可以翻译成这种语言，再由支持ONNX的推理引擎（如ONNX Runtime、TensorRT）读取并执行。

2.1 ONNX格式核心：计算图与节点

ONNX模型的核心是一个计算图（Graph）。图中包含三类基本元素：输入（Input）、输出（Output）和节点（Node）。节点代表算子（如Conv、Relu、MatMul），数据在节点间以张量（Tensor）形式流动。每个节点有明确的类型（op_type）、属性（attribute）和输入输出连接关系。

一个典型的ONNX模型文件（.onnx）内部采用Protobuf二进制格式存储，包含模型版本、生产者信息、计算图定义以及可选的元数据。下图展示了一个简化ONNX计算图的结构与数据流方向：

上图展示了一个典型的CNN推理子图。注意BatchNorm在推理时通常被折叠进前一个Conv的权重中（称为算子融合），这是ONNX导出工具自动完成的优化之一。理解这个结构对后续排查GPU推理性能瓶颈至关重要——例如，如果图中存在大量小算子（如逐元素Add），GPU的并行优势可能无法充分发挥。

2.2 模型导出：从PyTorch到ONNX

以PyTorch为例，导出ONNX使用torch.onnx.export函数。核心参数包括模型、示例输入、输出路径以及动态轴（dynamic_axes）设置。动态轴决定了导出后的模型是否允许批量大小（batch size）或图像尺寸变化。

import torch
import torchvision.models as models

model = models.resnet18(pretrained=True)
model.eval()  # 必须切换到eval模式，关闭Dropout和BatchNorm的训练行为

dummy_input = torch.randn(1, 3, 224, 224)

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"},   # 允许batch维度可变
        "output": {0: "batch_size"}
    }
)
print("导出完成")

2.3 模型加载与验证

导出后，必须验证模型的正确性和完整性。使用ONNX Runtime加载并执行一次推理是最直接的验证方式。同时，可以使用onnx.checker检查模型结构合法性。

import onnx
import onnxruntime as ort
import numpy as np

# 1. 结构检查
model = onnx.load("resnet18.onnx")
onnx.checker.check_model(model)
print("模型结构合法")

# 2. 推理验证
ort_session = ort.InferenceSession("resnet18.onnx", providers=["CPUExecutionProvider"])
input_name = ort_session.get_inputs()[0].name
output_name = ort_session.get_outputs()[0].name

# 构造与导出时一致的输入
test_input = np.random.randn(1, 3, 224, 224).astype(np.float32)
outputs = ort_session.run([output_name], {input_name: test_input})
print(f"输出shape: {outputs[0].shape}")

加载时需注意provider优先级。如果希望使用GPU，应显式指定CUDAExecutionProvider，否则默认使用CPU。验证阶段建议先用CPU跑通逻辑，再切换到GPU排查性能问题。

2.4 模型结构可视化

可视化是理解模型结构、定位导出问题的最有效手段。推荐以下工具：

使用Netron打开导出的resnet18.onnx，你应该能看到清晰的ResNet结构：卷积层、残差连接、全局平均池化和全连接层。如果发现以下异常，说明导出有问题：

• 出现大量Shape、Gather、Unsqueeze节点：通常源于动态shape操作，会降低GPU推理效率。
• BatchNorm节点仍然独立存在：说明算子融合未生效，可能是opset版本过低。
• 输入输出名称与预期不符：检查input_names和output_names设置。