torch.onnx¶
概述¶
开放神经网络交换 (ONNX) 是一种表示机器学习模型的开放标准格式。 torch.onnx 模块从原生 PyTorch torch.nn.Module 模型中捕获计算图并将其转换为 ONNX 图。
导出的模型可以被许多支持ONNX的运行时使用，包括 微软的ONNX 运行时。
你可以使用以下两种 ONNX 导出器 API：
基于 TorchDynamo 的 ONNX 导出器¶
基于 TorchDynamo 的 ONNX 导出器是 PyTorch 2.1 及更新版本的最新（且为 Beta 版）导出器。
利用 TorchDynamo 引擎挂钩到 Python 的帧评估 API，并将其字节码动态地重写为 FX 图。然后对生成的 FX 图进行优化处理，最后将其转换为 ONNX 图。
这种方法的主要优势在于，FX图是通过字节码分析捕获的，这种方式保留了模型的动态特性，而不是使用传统的静态跟踪技术。
基于 TorchScript 的 ONNX 导出器¶
基于 TorchScript 的 ONNX 导出器自 PyTorch 1.2.0 起可用。
TorchScript 用于追踪（通过 torch.jit.trace()）
模型并捕获静态计算图。
因此，生成的图有一些限制：
它不会记录任何控制流，例如 if 语句或循环；
不处理
training和eval模式之间的细微差别；不真正处理动态输入
为了支持静态追踪的限制，导出器还支持 TorchScript 脚本编写
(通过 torch.jit.script())，这增加了对数据依赖控制流的支持，例如。但是，TorchScript
本身是 Python 语言的一个子集，因此并非所有 Python 功能都得到支持，比如原地操作。