Shortcuts

torch.onnx

概述

开放神经网络交换 (ONNX) 是一种表示机器学习模型的开放标准格式。torch.onnx 模块从本地的 PyTorch torch.nn.Module 模型中捕获计算图,并将其转换为 ONNX 图

导出的模型可以被众多支持ONNX的运行时之一使用,包括微软的ONNX Runtime

有两种ONNX导出器API可供使用,如下所示:

基于TorchDynamo的ONNX导出器

基于TorchDynamo的ONNX导出器是PyTorch 2.0及更新版本中最新的(且为Beta版)导出器

TorchDynamo 引擎被用来钩入 Python 的帧评估 API,并动态地将字节码重写为 FX 图。生成的 FX 图在最终转换为 ONNX 图之前会进行优化。

这种方法的主要优点是,FX 图是通过字节码分析捕获的,保留了模型的动态特性,而不是使用传统的静态追踪技术。

了解更多关于基于TorchDynamo的ONNX导出器的信息

基于TorchScript的ONNX导出器

自 PyTorch 1.2.0 起,基于 TorchScript 的 ONNX 导出器可用

TorchScript 被用来追踪(通过 torch.jit.trace())模型并捕获一个静态计算图。

因此,生成的图表有一些局限性:

  • 它不会记录任何控制流,如if语句或循环;

  • 不处理训练评估模式之间的细微差别;

  • 不真正处理动态输入

为了支持静态追踪的局限性,导出器还支持通过torch.jit.script()进行TorchScript脚本编写,从而增加了对数据依赖控制流的支持,例如。然而,TorchScript本身是Python语言的一个子集,因此并非所有Python中的功能都受支持,例如就地操作。

了解更多关于基于TorchScript的ONNX导出器的信息

贡献 / 开发

ONNX 导出器是一个社区项目,我们欢迎贡献。我们遵循 PyTorch 贡献指南,但您可能 也对阅读我们的 开发维基感兴趣。

优云智算