ONNX 简介 || 将 PyTorch 模型导出到 ONNX || 扩展 ONNX 注册表

ONNX 简介

创建于：2023 年 10 月 4 日 | 最后更新：2025 年 1 月 27 日 | 最后验证：2024 年 11 月 5 日

作者：Ti-Tai Wang 和 Xavier Dupré

开放神经网络交换 (ONNX) 是一种用于表示机器学习模型的开放标准格式。torch.onnx 模块提供 API，用于从原生 PyTorch torch.nn.Module 模型捕获计算图，并将其转换为 ONNX 图。

导出的模型可以被许多 支持 ONNX 的运行时使用，包括 Microsoft 的 ONNX Runtime。

注意

目前，ONNX 导出器 API 有两种风格，但本教程将重点介绍 torch.onnx.dynamo_export。

TorchDynamo 引擎被用于hook到 Python 的帧评估 API，并动态地将其字节码重写为 FX 图。生成的 FX 图在最终转换为 ONNX 图之前会进行润色。

这种方法的主要优点是，FX 图是使用字节码分析捕获的，它保留了模型的动态性质，而不是使用传统的静态追踪技术。

依赖项

需要 PyTorch 2.1.0 或更高版本。

ONNX 导出器依赖于额外的 Python 包

• ONNX 标准库

• ONNX Script 库，使开发者能够以富有表现力且简单的方式使用 Python 子集编写 ONNX 运算符、函数和模型

• ONNX Runtime 加速机器学习库。

它们可以通过 pip 安装

pip install --upgrade onnx onnxscript onnxruntime

要验证安装，请运行以下命令

import torch
print(torch.__version__)

import onnxscript
print(onnxscript.__version__)

from onnxscript import opset18  # opset 18 is the latest (and only) supported version for now

import onnxruntime
print(onnxruntime.__version__)

每个 import 都必须成功且没有任何错误，并且必须打印库版本。

进一步阅读

下面的列表指的是从基本示例到高级场景的教程，不一定按照列出的顺序排列。您可以直接跳转到您感兴趣的特定主题，或者坐下来愉快地浏览所有教程，了解关于 ONNX 导出器的所有内容。

1. 将 PyTorch 模型导出到 ONNX

2. 扩展 ONNX 注册表