Lei Mao's Log Book
·
PyTorch自定义操作
内容提要
本文介绍了如何在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宏可以方便地注册这些操作，确保在不同设备上（如CPU和GPU）都能正确调用相应的实现。这种灵活性使得开发者能够在深度学习模型中集成高效的自定义功能。
模型导出与推理的兼容性
为了确保自定义操作在模型导出时的兼容性，需要注册“假”的自定义类和函数版本。这一过程对于使用torch.export和torch.compile至关重要，因为它允许在符号跟踪时不执行实际的C++/CUDA代码，从而提高了模型的可移植性和灵活性。
C++推理程序的优势
通过使用dlopen加载自定义操作库，C++推理程序可以在没有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实现注册自定义操作。
