ONNX简介 || 将PyTorch模型导出为ONNX || 扩展ONNX导出器操作符支持 || 将带控制流的模型导出为ONNX

ONNX简介

Created On: Oct 04, 2023 | Last Updated: Mar 05, 2025 | Last Verified: Nov 05, 2024

作者: Ti-Tai Wang, Thiago Crepaldi。

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

导出的模型可以被任何支持ONNX的`运行时 <https://onnx.ai/supported-tools.html#deployModel>`_ 消耗，包括微软的`ONNX Runtime <https://www.onnxruntime.ai>`_。

备注

目前，您可以通过`TorchScript https://pytorch.org/docs/stable/jit.html`_ 或 `ExportedProgram https://pytorch.org/docs/stable/export.html`_ 之一根据`torch.onnx.export <https://pytorch.org/docs/stable/onnx_torchscript.html#torch.onnx.export>`_的布尔参数dynamo将模型导出为ONNX。在本教程中，我们将重点介绍``ExportedProgram``方法。

设置``dynamo=True``时，导出器将使用`torch.export <https://pytorch.org/docs/stable/export.html>`_ 捕获``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. Exporting a PyTorch model to ONNX

2. Extending the ONNX exporter operator support

3. Export a model with control flow to ONNX