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

ONNX 简介

作者：Thiago Crepaldi,

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

导出的模型可以被许多支持 ONNX 的 运行时 中的任何一个使用，包括 Microsoft 的 ONNX 运行时.

注意

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

TorchDynamo 引擎用于挂钩到 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 注册表