Lei Mao's Log Book
·
PyTorch 导出
💡
原文英文，约800词，阅读约需3分钟。
📝
内容提要
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中保存和加载导出的模型？
可以使用torch.export.save保存模型，使用torch.export.load加载模型。
导出的模型在C++中如何运行？
导出的模型必须使用AOTInductor或Executorch进行优化和编译，才能在C++中运行。
PyTorch导出模型与ONNX模型有什么相似之处？
PyTorch导出模型类似于ONNX模型，PyTorch IR类似于ONNX Opset。
➡️
