ONNX 简介 || 将 PyTorch 模型导出到 ONNX || 扩展 ONNX 导出器算子支持 || 将带有控制流的模型导出到 ONNX

ONNX 简介

创建日期：2023 年 10 月 04 日 | 最后更新：2025 年 03 月 05 日 | 最后验证：2024 年 11 月 05 日

作者：Ti-Tai Wang, Thiago Crepaldi。

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

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

注意

目前，您可以通过 `TorchScript https://pytorch.ac.cn/docs/stable/jit.html`_ 或 `ExportedProgram https://pytorch.ac.cn/docs/stable/export.html`_ 来选择将模型导出到 ONNX，通过 torch.onnx.export 中的布尔参数 dynamo 进行控制。在本教程中，我们将重点介绍 ExportedProgram 方法。

当设置 dynamo=True 时，导出器将使用 torch.export 捕获一个 ExportedProgram，然后将图转换为 ONNX 表示。这种方法是将模型导出到 ONNX 的新推荐方式。它与 PyTorch 2.0 特性配合更稳健，对较新的 ONNX 算子集有更好的支持，并且消耗的资源更少，使得导出大型模型成为可能。

依赖项

需要 PyTorch 2.5.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__)

import onnxruntime
print(onnxruntime.__version__)

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

延伸阅读

以下列表包含从基本示例到高级场景的教程，不一定按列出的顺序排列。请随意直接跳到您感兴趣的特定主题，或者耐心仔细地学习所有教程，掌握 ONNX 导出器的所有知识。

1. 将 PyTorch 模型导出到 ONNX

2. 扩展 ONNX 导出器算子支持

3. 将带有控制流的模型导出到 ONNX