注意
点击 here 下载完整的示例代码
ONNX简介 || 将PyTorch模型导出到ONNX || 扩展ONNX注册表
ONNX简介
创建于:2023年10月04日 | 最后更新:2024年9月05日 | 最后验证:2024年11月05日
作者: Thiago Crepaldi,
Open Neural Network eXchange (ONNX) 是一个用于表示机器学习模型的开放标准格式。torch.onnx
模块提供了API,可以从原生的 PyTorch torch.nn.Module
模型中捕获计算图并将其转换为 ONNX 图。
导出的模型可以被任何支持ONNX的运行时使用,包括微软的ONNX Runtime。
注意
目前,有两种类型的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必须成功且没有任何错误,并且必须打印出库的版本。