ONNX 教程

欢迎学习 ONNX (Open Neural Network Exchange)！本教程将带你系统掌握 ONNX 的核心概念、模型导出、推理部署以及性能优化的完整流程。

什么是 ONNX？

ONNX（开放神经网络交换）是一种开放的模型表示格式，由微软、亚马逊、Meta（原 Facebook）等公司于 2017 年联合推出并开源。它的核心目标是解决深度学习领域长期存在的"框架锁定"问题——在不同训练框架和推理引擎之间建立一座标准化的桥梁。

深度学习的发展催生了众多优秀的框架：PyTorch 以其动态图和易用性在研究领域占据主导地位，TensorFlow 则在生产部署方面积累了大量用户。然而，用 PyTorch 训练的模型想要部署到 TensorRT 上加速，或者 TensorFlow 模型想要在 ONNX Runtime 中运行，都需要额外的转换工作。ONNX 正是为了解决这个问题而生。

ONNX 定义了一套标准化的算子集合和数据格式，任何支持 ONNX 的框架都可以读取和执行模型，无需关心模型最初是在哪里训练的。这就像是深度学习领域的"通用语言"。

ONNX 的核心优势

特性	说明
框架互操作性	在 PyTorch 中训练，在 TensorFlow 中推理，或使用 TensorRT 加速，一次导出，处处可用
硬件广泛支持	NVIDIA TensorRT、Intel OpenVINO、AMD ROCm、华为昇腾等主流硬件平台均原生支持 ONNX
版本兼容性	通过 Opset 机制保证模型的前向兼容性，旧版本模型在新版本推理引擎中仍可运行
易于优化	标准化的图结构使得模型优化工具（如 onnx-simplifier、ONNX GraphSurgeon）能够高效工作
生产级稳定性	ONNX Runtime 作为微软官方维护的推理引擎，已被 Windows、Azure 等产品大规模验证

ONNX 模型文件结构

一个 .onnx 文件本质上是一个 Protocol Buffers 序列化文件，包含以下核心组件：

ModelProto：模型的根节点，包含版本信息、生产者信息等元数据
GraphProto：定义计算图的结构，包括节点、边、输入输出
NodeProto：每个计算节点，包含算子类型、输入输出名称、属性等
TensorProto：存储模型的权重参数
OpsetImport：声明模型使用的算子集版本

ONNX 生态概览

训练框架支持

框架	导出方式	支持程度
PyTorch	`torch.onnx.export()`	官方支持，最完善
TensorFlow/Keras	tf2onnx	社区维护，覆盖广泛
Scikit-learn	sklearn-onnx	传统机器学习模型支持
MXNet	mxnet.contrib.onnx	官方支持
JAX	jax2onnx	实验性支持

推理引擎选择

引擎	特点	适用场景
ONNX Runtime	跨平台、易用、官方维护	通用推理、快速部署
TensorRT	NVIDIA GPU 专属、极致性能	实时推理、大模型部署
OpenVINO	Intel 硬件优化	Intel CPU/GPU 部署
Triton Inference Server	多模型服务、动态批处理	生产环境、微服务架构

为什么学习 ONNX？

解决生产环境的核心痛点

在实际的 AI 工程中，训练和推理往往使用不同的工具链：

训练阶段：选择 PyTorch，因为它提供了灵活的调试能力和活跃的学术社区支持
推理阶段：选择 C++、TensorRT 或 ONNX Runtime，因为它们提供了更高的性能和更低的资源占用

这两个阶段之间存在巨大的鸿沟。如果没有 ONNX，你需要：

在 PyTorch 中训练模型
手动将模型权重导出为某种中间格式
在推理端重新构建模型结构
加载权重并验证结果一致性

这个过程不仅繁琐，而且容易出错。ONNX 的出现让这个过程变得简单可靠。

实际应用案例

案例一：模型服务化部署

某电商平台的商品识别模型最初在 PyTorch 中开发和训练。上线时，团队选择将模型导出为 ONNX 格式，部署在 ONNX Runtime 上。这样做的好处是：

推理服务不需要安装 PyTorch，体积从 2GB 缩减到 200MB
Python 推理延迟从 50ms 降低到 15ms
可以轻松切换到 TensorRT 获得更快的速度

案例二：多硬件平台适配

某自动驾驶公司需要将同一个感知模型部署到多种硬件：车载 NVIDIA Orin、边缘计算盒子、云端服务器。通过 ONNX 作为中间格式，他们只需要维护一个 .onnx 文件，各硬件平台上的推理引擎都能直接加载运行。

案例三：模型优化流水线

在模型上线前，通常会经历这样的流程：

每一步都有成熟的工具支持，形成标准化的工业流程。

核心概念：Opset 版本

ONNX 的算子集（Operator Set，简称 Opset）是理解 ONNX 的关键概念。每个 Opset 版本定义了一组标准算子及其行为规范。

Opset 版本演进

Opset	发布时间	主要更新
11	2019	基础算子完善，广泛兼容
12	2019	NonMaxSuppression 改进
13	2020	更好的类型支持
14	2020	CumSum、Trilu 等新算子
15	2021	ScatterND、GatherND 增强
17	2022	LayerNormalization 等新算子
18-21	2023-2024	持续扩展新算子

如何选择 Opset 版本

选择 Opset 版本时需要考虑两个因素：

推理引擎的兼容性：目标推理引擎支持的 Opset 最高版本是多少？TensorRT 8.x 支持 Opset 17，而 ONNX Runtime 1.15 支持到 Opset 20。
模型算子需求：你的模型是否使用了某个 Opset 才引入的新算子？例如，如果模型使用了 LayerNormalization 算子，需要至少 Opset 17。

# 导出时指定 Opset 版本
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    opset_version=17  # 根据目标平台选择
)

实际建议是：在满足功能需求的前提下，选择较低的 Opset 版本以获得更广泛的兼容性。对于大多数视觉模型，Opset 14-17 是安全的选择。

环境安装

安装核心库

# 安装 ONNX 核心库和 ONNX Runtime
pip install onnx onnxruntime

# 安装模型简化工具
pip install onnx-simplifier

# 如果需要 GPU 推理
pip install onnxruntime-gpu

验证安装

import onnx
import onnxruntime as ort

# 检查版本
print(f"ONNX 版本: {onnx.__version__}")
print(f"ONNX Runtime 版本: {ort.__version__}")

# 检查可用的执行提供者（Execution Providers）
print(f"可用的执行提供者: {ort.get_available_providers()}")
# 输出示例: ['CPUExecutionProvider', 'CUDAExecutionProvider']

# 检查当前设备
print(f"当前推理设备: {ort.get_device()}")

常见安装问题

问题一：onnx 和 onnxruntime 版本不兼容

ONNX 和 ONNX Runtime 有版本依赖关系。如果遇到 onnx.onnx_cpp2py_export.checker.ValidationError 等错误，可能是版本不匹配。建议安装兼容版本：

# 推荐的版本组合
pip install onnx==1.15.0 onnxruntime==1.16.3

问题二：GPU 版本无法识别 CUDA

安装 onnxruntime-gpu 后仍然无法使用 GPU，可能是因为 CUDA 版本不匹配。ONNX Runtime GPU 版本对 CUDA 有特定要求：

ONNX Runtime GPU	CUDA	cuDNN
1.16.x	11.8 / 12.x	8.x
1.15.x	11.8	8.x

检查 CUDA 是否正确安装：

import onnxruntime as ort

# 检查 CUDA 执行提供者是否可用
if 'CUDAExecutionProvider' in ort.get_available_providers():
    print("GPU 推理可用")
else:
    print("GPU 推理不可用，请检查 CUDA 安装")

教程目录

本教程按照实际工作流程组织，从模型导出到部署优化：

核心操作

PyTorch 导出 ONNX - 详解 torch.onnx.export 的每个参数，处理动态维度和复杂模型
TensorFlow 导出 ONNX - 使用 tf2onnx 转换 TensorFlow/Keras 模型
Python 推理 - ONNX Runtime Python API 完整指南，包括 IOBinding 和异步推理
C++ 部署 - 工业级 C++ 推理代码，涵盖内存管理和多线程

进阶主题

模型优化 - 图优化、量化、剪枝等性能提升技术
ONNX 算子集 - 常见算子详解与自定义算子开发

参考资料

知识速查表 - 常用 API 和命令速查

学习建议

ONNX 看起来简单——不就是导出一个文件吗？但要在生产环境中用好它，需要理解很多细节：

先跑通基础流程：用最简单的模型完成导出-验证-推理的完整流程
理解 Opset 机制：这是解决兼容性问题的关键
掌握调试技巧：模型输出不一致时，知道如何定位问题
熟悉优化工具：onnx-simplifier、量化工具是性能调优的利器

接下来的章节将逐一展开这些主题。让我们从 PyTorch 模型导出开始。

什么是 ONNX？​

ONNX 的核心优势​

ONNX 模型文件结构​

ONNX 生态概览​

训练框架支持​

推理引擎选择​

为什么学习 ONNX？​

解决生产环境的核心痛点​

实际应用案例​

核心概念：Opset 版本​

Opset 版本演进​

如何选择 Opset 版本​

环境安装​

安装核心库​

验证安装​

常见安装问题​

教程目录​

核心操作​

进阶主题​

参考资料​

学习建议​