Lei Mao's Log Book
·
PyTorch自定义操作
💡
原文英文，约2900词，阅读约需11分钟。
📝
内容提要
本文介绍了如何在PyTorch中实现自定义操作，包括使用C++和CUDA编写自定义函数和类。通过示例代码，展示了如何注册和使用这些自定义操作，确保它们在PyTorch模型和AOTInductor编译的推理程序中正常工作，并讨论了模型导出及推理验证。
🎯
关键要点
- PyTorch自定义操作可以通过C++和CUDA实现，包括自定义类和函数。
- 使用TORCH_LIBRARY_IMPL宏注册自定义函数，支持CPU和CUDA实现。
- 自定义类可以使用torch::CustomClassHolder定义，并通过TORCH_LIBRARY宏注册。
- 自定义操作和类的实现被编译为共享库，可以通过torch.ops.load_library加载。
- 为了支持torch.export和torch.compile，需要注册“假”的自定义类和函数版本。
- 导出的模型可以使用torch._inductor.aoti_compile_and_package编译并打包。
- 在C++推理程序中，可以使用dlopen加载自定义操作库，无需pybind11或libpython依赖。
❓
延伸问答
如何在PyTorch中实现自定义操作？
可以通过C++和CUDA实现自定义操作，包括自定义类和函数，并使用TORCH_LIBRARY_IMPL宏进行注册。
自定义操作如何在PyTorch模型中使用？
自定义操作可以通过torch.ops.load_library加载，并在PyTorch模型中调用。
如何确保自定义操作在模型导出时正常工作？
需要注册“假”的自定义类和函数版本，以支持torch.export和torch.compile。
自定义类在PyTorch中如何定义和注册？
可以使用torch::CustomClassHolder定义自定义类，并通过TORCH_LIBRARY宏进行注册。
如何在C++推理程序中加载自定义操作库？
可以使用dlopen加载自定义操作库，无需pybind11或libpython依赖。
自定义操作的CPU和CUDA实现如何注册？
使用TORCH_LIBRARY_IMPL宏分别为CPU和CUDA实现注册自定义操作。
➡️
