Lei Mao's Log Book
·
PyTorch 导出
内容提要
PyTorch 2引入torch.export功能，允许将模型导出为静态图以优化推理性能。自定义操作可通过无状态的torch.ops或有状态的torch.classes定义。导出时需避免数据依赖形状和控制流，以确保模型可静态表示。导出的模型可在Python中保存和加载，但在C++中需使用AOTInductor或Executorch进行优化和编译，确保模型无图断裂是推理的关键。
关键要点
-
PyTorch 2引入torch.export功能，允许将模型导出为静态图以优化推理性能。
-
自定义操作可以通过无状态的torch.ops或有状态的torch.classes定义。
-
导出时需避免数据依赖形状和控制流，以确保模型可静态表示。
-
导出的模型可以在Python中保存和加载，但在C++中需使用AOTInductor或Executorch进行优化和编译。
-
确保模型无图断裂是推理的关键。
-
自定义操作的实现可以使用假Python类来追踪并包含在导出的图中。
-
导出模型不能在C++中加载和运行，必须先进行优化和编译。
-
PyTorch导出模型类似于ONNX模型，PyTorch IR类似于ONNX Opset。
延伸解读
导出模型的关键注意事项
在使用torch.export导出模型时，确保模型中没有数据依赖的形状和控制流是至关重要的。这意味着在设计模型时，开发者需要避免使用动态变化的输入，这样才能确保模型能够被静态表示并顺利导出。
自定义操作的实现方式
在PyTorch中实现自定义操作时，推荐使用无状态的torch.ops函数进行定义。这种方式不仅简化了实现过程，还能确保自定义操作能够正确地被导出到C++中。相对而言，有状态的torch.classes实现则较为复杂，需在C++和Python中都进行定义。
导出模型的运行环境
导出的模型可以在Python中直接保存和加载，但在C++中使用时必须经过优化和编译。这意味着开发者在选择部署环境时，需要考虑到模型的导出和运行流程，以确保推理性能的最优化。
延伸问答
PyTorch 2的torch.export功能有什么作用？
torch.export功能允许将PyTorch模型导出为静态图，以优化推理性能。
如何定义自定义操作以便于导出？
可以通过无状态的torch.ops函数或有状态的torch.classes类来定义自定义操作，推荐使用无状态的方式。
导出模型时需要避免哪些情况？
导出时需避免数据依赖形状和控制流，以确保模型可静态表示。
如何在Python中保存和加载导出的模型？
可以使用torch.export.save保存模型，使用torch.export.load加载模型。
导出的模型在C++中如何运行？
导出的模型必须使用AOTInductor或Executorch进行优化和编译，才能在C++中运行。
PyTorch导出模型与ONNX模型有什么相似之处？
PyTorch导出模型类似于ONNX模型，PyTorch IR类似于ONNX Opset。
