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。
延伸问答
PyTorch 2中的torch.export功能有什么作用？
torch.export功能允许将PyTorch模型导出为静态图，以优化推理性能。
如何定义自定义操作以便于导出？
自定义操作可以通过无状态的torch.ops函数或有状态的torch.classes类来定义，推荐使用无状态的方式。
导出模型时需要避免哪些问题？
导出时需避免数据依赖形状和控制流，以确保模型可静态表示。
导出的模型如何在Python和C++中使用？
导出的模型可以在Python中使用torch.export.save和torch.export.load进行保存和加载，但在C++中需使用AOTInductor或Executorch进行优化和编译。
为什么导出模型不能在C++中直接加载和运行？
导出的模型必须先进行优化和编译，才能在C++中运行，直接加载和运行是不支持的。
PyTorch导出模型与ONNX模型有什么相似之处？
PyTorch导出模型类似于ONNX模型，PyTorch IR类似于ONNX Opset，二者都不允许数据依赖形状和控制流。
