2、ONNX基础：ONNX格式介绍、模型导出流程、常见工具链

2.1 ONNX格式核心概念

ONNX（Open Neural Network Exchange）是一个开放的模型表示标准，用于在不同深度学习框架之间交换模型。它由微软和Facebook于2017年联合发起，目前由Linux基金会旗下的ONNX社区维护。理解ONNX格式，需要先掌握三个核心对象：计算图（Graph）、节点（Node）和张量（Tensor）。

ONNX模型本质上是一个有向无环图（DAG）。图中的每个节点代表一个算子（如Conv、Relu、MatMul），边则代表张量数据在算子之间的流动。模型文件采用Protobuf序列化存储，扩展名通常为.onnx。一个完整的ONNX模型包含以下关键信息：

  • 模型元信息：ir_version（IR版本）、opset_import（算子集版本）、producer_name（导出框架名称）等。
  • 计算图定义：graph节点，包含所有算子节点、初始权重（initializer）和输入输出张量。
  • 输入输出声明：每个输入/输出张量的名称、数据类型（如float32、int64）和维度信息（支持动态维度）。
  • 属性（Attribute）：算子的静态参数，如卷积的kernel_shape、strides、pads等。
关键理解：ONNX不是“另一种深度学习框架”，它不负责训练或推理，只负责定义一种中间表示。你可以把ONNX理解为“模型界的通用语言”——PyTorch模型翻译成ONNX，TensorFlow模型也翻译成ONNX，下游的推理引擎（如ONNX Runtime、TensorRT）再统一从ONNX读取模型并执行。

下面这张图展示了ONNX模型内部的结构层次，以及它与框架、推理引擎之间的关系：

PyTorch torch.onnx.export TensorFlow tf2onnx.convert PaddlePaddle paddle2onnx ONNX 模型文件 元信息 (ir_version, opset) Graph 计算图 Node (Conv) Node (Relu) Initializer 权重 ONNX Runtime CPU / CUDA 推理 TensorRT NVIDIA GPU 加速 OpenVINO Intel CPU / VPU 模型导出 = 将框架模型“翻译”为ONNX计算图 导出报错通常发生在“翻译”过程中：算子不支持、属性不兼容、动态维度未声明等

2.2 模型导出流程

从PyTorch/TensorFlow模型到可运行的ONNX文件，需要经过一条明确的流水线。以最常见的PyTorch导出为例，完整流程如下：

  1. 准备模型与输入：将模型切换到eval()模式，并构造一个符合输入维度的示例张量（dummy input）。
  2. 调用导出API：执行torch.onnx.export()，传入模型、示例输入、输出路径和配置参数。
  3. 执行符号跟踪（Tracing）：PyTorch会实际运行一次模型前向传播，记录所有张量操作，将其转换为ONNX算子序列。
  4. 算子映射与转换：PyTorch算子逐一映射到ONNX算子。若遇到无法映射的算子，会抛出UnsupportedOperatorError。
  5. 常量折叠与优化：部分固定计算（如形状推导、常量运算）会被提前计算并写入模型，减少运行时开销。
  6. 生成并校验ONNX文件：写入.onnx文件，可使用onnx.checker或onnxruntime加载验证。

一个典型的PyTorch导出代码如下：

import torch
import torchvision.models as models

# 1. 准备模型与输入
model = models.resnet18(pretrained=True)
model.eval()
dummy_input = torch.randn(1, 3, 224, 224)

# 2. 调用导出API
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("导出完成")

导出完成后，建议立即用ONNX Runtime做一次加载和推理验证，确保模型文件可用：

import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession("resnet18.onnx")
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name
result = sess.run([output_name], {input_name: np.random.randn(1, 3, 224, 224).astype(np.float32)})
print(f"推理输出shape: {result[0].shape}")

2.3 常见工具链全景

围绕ONNX生态，已经形成了一套完整的工具链。下表总结了开发中最常用的工具及其用途：

工具名称 主要用途 典型场景
torch.onnx PyTorch → ONNX导出 PyTorch模型转换
tf2onnx TensorFlow → ONNX转换 TF SavedModel / Keras模型转换
paddle2onnx PaddlePaddle → ONNX转换 飞桨模型转换
onnx (Python包) ONNX模型解析、修改、检查 查看节点、修改权重、模型合并
onnxruntime ONNX模型推理引擎 跨平台部署、性能验证
onnxsim ONNX模型简化与优化 删除冗余节点、常量折叠
netron 模型可视化工具 图形化查看计算图结构
onnx2trt ONNX → TensorRT引擎转换 NVIDIA GPU高性能推理

2.4 典型场景与误区复盘

场景一：动态输入尺寸的模型导出。目标检测模型（如YOLO）的输入图片尺寸可能不固定。若在导出时未声明dynamic_axes，导出的模型会固定输入shape为[1, 3, 640, 640]，后续传入其他尺寸图片会直接报错。正确做法是在导出时指定动态维度，并在推理时通过sess.get_inputs()[0].shape确认动态轴是否生效。

场景二：自定义算子的导出失败。当模型中含有PyTorch原生算子无法直接映射到ONNX时（如某些自定义autograd.Function），导出会抛出UnsupportedOperatorError。此时有两种解决思路：一是用PyTorch原生算子重写自定义逻辑；二是通过torch.onnx.register_custom_op_symbolic注册自定义映射函数。

常见误区：很多开发者认为“模型能导出成功，就一定能正确推理”。实际上，导出成功只代表计算图结构转换完成，但可能存在以下隐患：
  • 算子精度差异：某些算子在ONNX Runtime中的实现与PyTorch存在微小数值差异，可能导致输出不一致。
  • 动态维度未生效：导出时虽然声明了动态轴，但模型内部某些算子（如Reshape）仍可能隐式固定了维度。
  • opset版本过低：旧版opset可能缺少新算子支持，导致导出时自动替换为多个子算子组合，增加图复杂度和推理延迟。
建议每次导出后，用相同的输入分别在PyTorch和ONNX Runtime中运行，对比输出差异（通常要求误差在1e-4以内）。

场景三：模型简化与优化。导出的ONNX模型往往包含大量冗余的恒等映射、形状转换节点。使用onnxsim可以自动简化模型：

import onnx
from onnxsim import simplify

model = onnx.load("resnet18.onnx")
model_simp, check = simplify(model)
assert check, "简化后模型校验失败"
onnx.save(model_simp, "resnet18_sim.onnx")
print(f"节点数: {len(model.graph.node)} → {len(model_simp.graph.node)}")
← 上一章 📖 返回目录 下一章 →