注意
点击 这里 下载完整示例代码
ONNX简介 || 将PyTorch模型导出到ONNX || 扩展ONNX注册表
ONNX简介¶
创建日期: 2023年10月4日 | 最后更新日期: 2024年9月5日 | 最后验证日期: 2024年11月5日
作者： 蒂亚戈·克雷帕尔迪，
开放神经网络交换 (ONNX) 是一种用于表示机器学习模型的开放标准格式。 torch.onnx 模块提供了从原生 PyTorch torch.nn.Module 模型中捕获计算图并将其转换为 ONNX 图的 API。
导出的模型可以被许多支持ONNX的运行时使用， 包括微软的ONNX 运行时。
注意
目前，有两类ONNX导出API，
但本教程将专注于torch.onnx.dynamo_export。
TorchDynamo 引擎被用来钩入 Python 的帧评估 API，并动态地将其字节码重写为 FX 图。 生成的 FX 图在最终转换为 ONNX 图 之前进行了优化。
这种方法的主要优势在于，FX图是通过字节码分析捕获的，这种方式保留了模型的动态特性，而不是使用传统的静态跟踪技术。
依赖项¶
需要 PyTorch 2.1.0 或更高版本。
ONNX导出器依赖于额外的Python包：
ONNX standard library
ONNX Script library that enables developers to author ONNX operators, functions and models using a subset of Python in an expressive, and yet simple fashion
ONNX Runtime accelerated machine learning library.
它们可以通过 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必须成功运行且不出现任何错误，并且库版本必须被打印出来。