2. ONNX基础

2.1 ONNX是什么：不止是“模型文件”

ONNX（Open Neural Network Exchange）本质上是一个开放、跨框架的模型中间表示标准。它不负责训练，也不直接执行推理，而是定义了一套统一的计算图描述格式，让模型能够在不同框架（PyTorch、TensorFlow、PaddlePaddle等）和不同推理引擎（ONNX Runtime、TensorRT、OpenVINO）之间自由流转。你可以把它理解为深度学习的“通用语言”——训练框架负责“写诗”，ONNX负责“翻译”，推理引擎负责“朗诵”。

对于CPU推理优化而言，ONNX的价值在于：它剥离了训练框架的运行时开销，将模型固化为静态的计算图，从而让推理引擎能够对图结构进行深度优化（算子融合、内存规划、指令调度等）。这也是为什么在工业部署中，ONNX几乎成为模型交换的事实标准。

核心认知：ONNX不是一个新的深度学习框架，而是一座连接“训练生态”与“部署生态”的桥梁。它的存在让“训练一次，到处推理”成为可能。

2.2 ONNX模型结构：计算图、节点与张量

一个完整的ONNX模型文件（通常以.onnx为后缀）内部是一个序列化的Protobuf结构，其核心逻辑模型由以下几部分组成：

  • 计算图（Graph）：模型的主体，描述数据从输入到输出的完整流动过程。
  • 节点（Node）：图中的每个算子（Conv、Relu、MatMul等）对应一个节点，节点包含算子类型（op_type）、输入输出张量名、属性（attribute，如卷积的kernel_shape）。
  • 张量（Tensor）：数据以张量形式在节点间流动，每个张量有名称、数据类型（float32、int64等）和形状（shape）。
  • 初始器（Initializer）：存储模型的权重和偏置，它们作为常量输入挂在图上，不随推理过程改变。
  • 元数据（Metadata）：包括模型版本、算子集版本（opset）、生产者信息等，用于兼容性判断。

下面这张图展示了一个简化ONNX模型的结构，帮助你直观理解数据如何从输入经过多个算子节点流向输出：

ONNX 模型内部结构示意 输入张量 data_0 [1,3,224,224] Conv kernel_shape=[3,3] strides=[1,1] Relu 激活函数 MaxPool kernel_shape=[2,2] strides=[2,2] 输出张量 prob [1,1000] 初始器 conv_w [64,3,3,3] 初始器 conv_b [64] data_0 conv_out relu_out pool_out 权重 偏置 输入/输出 算子节点 初始器（权重） 数据流 权重绑定

从图中可以看到，ONNX模型本质上是一个有向无环图（DAG）。数据从输入张量出发，依次经过Conv、Relu、MaxPool等节点，每个节点消费上游张量并产生下游张量，最终汇聚到输出。初始器（权重）不参与数据流动的主线，而是作为“旁路”挂载到对应算子节点上。这种结构清晰、无循环依赖的特点，使得推理引擎可以方便地对整个图进行静态分析和优化。

2.3 导出ONNX模型：PyTorch与TensorFlow实战

不同框架导出ONNX的方式略有差异，但核心思路一致：将动态的训练图“固化”为静态的推理图。下面以最常见的PyTorch为例，展示标准导出流程：

import torch
import torchvision.models as models

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

# 2. 构造一个示例输入（batch=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("ONNX模型导出成功！")

对于TensorFlow（Keras）模型，导出方式如下：

import tensorflow as tf
import tf2onnx

# 加载Keras模型
model = tf.keras.applications.MobileNetV2(weights="imagenet")
model.compile()

# 转换为ONNX
spec = (tf.TensorSpec((1, 224, 224, 3), tf.float32, name="input"),)
onnx_model = tf2onnx.convert.from_keras(model, input_signature=spec, opset=17)

# 保存为.onnx文件
with open("mobilenetv2.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())
print("转换完成！")
导出要点：导出前务必调用model.eval()（PyTorch）或确保模型处于推理模式，否则BatchNorm、Dropout等层会因训练/推理行为差异导致导出结果错误。同时，opset_version建议选择推理引擎支持的版本，ONNX Runtime通常支持较新的opset。

2.4 加载与验证ONNX模型

导出ONNX后，第一步是使用官方工具验证模型的完整性，然后加载到推理引擎中。ONNX Runtime是最常用的跨平台推理引擎，加载方式非常简洁：

import onnx
import onnxruntime as ort
import numpy as np

# 1. 验证模型结构
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)
print("模型结构验证通过！")

# 2. 创建推理会话
# 指定CPUExecutionProvider，确保使用CPU推理
session = ort.InferenceSession(
    "resnet18.onnx",
    providers=["CPUExecutionProvider"]
)

# 3. 查看输入输出信息
for inp in session.get_inputs():
    print(f"输入: {inp.name}, 形状: {inp.shape}, 类型: {inp.type}")
for out in session.get_outputs():
    print(f"输出: {out.name}, 形状: {out.shape}, 类型: {out.type}")

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

加载ONNX模型时，providers参数决定了使用哪个执行后端。在CPU环境下，指定CPUExecutionProvider即可；如果系统装有NVIDIA GPU，也可以添加CUDAExecutionProvider，ONNX Runtime会自动选择最优设备。

2.5 常见误区与复盘

误区后果正确做法
忘记调用model.eval()BatchNorm统计量使用训练值，推理结果错误导出前显式切换到评估模式
使用动态输入形状但不声明dynamic_axes模型只能接受固定尺寸输入，部署灵活性差在dynamic_axes中声明可变维度
opset版本设置过低新算子无法导出，或导出后算子集不兼容根据推理引擎支持范围选择opset（通常≥15）
忽略do_constant_folding模型包含冗余计算节点，推理效率降低导出时开启常量折叠，简化计算图
加载时不指定providers可能默认使用CPU，但未显式声明不利于性能调优明确指定CPUExecutionProvider或GPU Provider
重要提醒：ONNX模型导出后并非“一劳永逸”。不同推理引擎对ONNX算子的支持程度不同，某些自定义算子（如自定义ROIAlign）可能无法直接转换。遇到这种情况，需要检查算子是否在ONNX算子集内，或者考虑使用onnx-simplifier等工具对模型进行简化和兼容性处理。

掌握ONNX模型的结构与导入导出流程，是后续进行CPU推理优化的基础。只有理解了模型在ONNX格式下的组织方式，才能更有效地利用ONNX Runtime的图优化、算子融合、线程调度等高级特性，为模型在CPU上的高效运行打下坚实基础。

← 上一章 📖 返回目录 下一章 →