PyTorch 自定义操作符

创建于：2024年6月18日 | 最后更新：2024年7月22日 | 最后验证：2024年11月5日

PyTorch 提供了一个庞大的运算符库，这些运算符作用于张量（例如 torch.add， torch.sum 等）。然而，您可能希望将一个新的自定义操作引入 PyTorch， 并使其与 torch.compile、autograd 和 torch.vmap 等子系统一起工作。 为此，您必须通过 Python 的 torch.library 文档 或 C++ 的 TORCH_LIBRARY API 向 PyTorch 注册自定义操作。

从Python编写自定义操作符

请参阅自定义Python操作符。

如果您希望从Python（而不是C++）编写自定义操作符，您可能希望：

• 你有一个Python函数，你希望PyTorch将其视为一个不透明的可调用对象，特别是在torch.compile和torch.export方面。

• 你有一些Python绑定到C++/CUDA内核，并希望这些能与PyTorch子系统（如torch.compile或torch.autograd）组合使用

将自定义C++和/或CUDA代码与PyTorch集成

请参阅自定义C++和CUDA运算符。

如果您希望从C++（而不是Python）编写自定义操作符，您可能希望：

• 你有自定义的C++和/或CUDA代码。

• 你计划使用此代码与AOTInductor进行无Python推理。

自定义操作符手册

有关教程和本页面未涵盖的信息，请参阅 The Custom Operators Manual （我们正在努力将信息迁移到我们的文档站点）。我们建议您 首先阅读上述教程之一，然后将自定义操作符手册作为参考； 它并不适合从头到尾阅读。

何时应该创建自定义操作符？

如果你的操作可以表示为内置PyTorch操作符的组合，那么请将其编写为Python函数并调用它，而不是创建一个自定义操作符。如果你正在调用一些PyTorch无法理解的库（例如自定义C/C++代码、自定义CUDA内核或C/C++/CUDA扩展的Python绑定），请使用操作符注册API来创建自定义操作符。

为什么我应该创建一个自定义操作符？

可以通过获取Tensor的数据指针并将其传递给pybind的内核来使用C/C++/CUDA内核。然而，这种方法无法与PyTorch子系统（如autograd、torch.compile、vmap等）组合使用。为了使操作能够与PyTorch子系统组合使用，必须通过操作符注册API进行注册。