注意
跳转至页面底部 下载完整示例代码。
ONNX 简介 || 将 PyTorch 模型导出为 ONNX || 扩展 ONNX 导出器的算子支持 || 将带有控制流的模型导出为 ONNX
ONNX 简介#
创建日期:2023年10月4日 | 最后更新:2026年4月10日 | 最后验证:2024年11月5日
作者:Ti-Tai Wang, Thiago Crepaldi。
开放神经网络交换 (ONNX) 是一种用于表示机器学习模型的开放标准格式。torch.onnx 模块提供了一系列 API,用于捕获原生 PyTorch torch.nn.Module 模型的计算图,并将其转换为 ONNX 图。
导出的模型可以被许多支持 ONNX 的运行时所使用,包括微软的 ONNX Runtime。
当设置 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 导出器所有知识的乐趣。