PyTorch 自定义运算符¶
创建时间： Jun 18， 2024 |上次更新时间： 2024 年 7 月 22 日 |上次验证： Nov 05， 2024
PyTorch 提供了一个大型运算符库，这些运算符适用于张量（例如 、 、 等）。但是，您可能希望为 PyTorch 引入新的自定义操作
并使其与 、autograd 和 等子系统一起使用。
为此，您必须通过 Python torch.library 文档或 C++ API 向 PyTorch 注册自定义操作。torch.addtorch.sumtorch.compiletorch.vmapTORCH_LIBRARY
从 Python 编写自定义运算符¶
请参阅自定义 Python 运算符。
如果满足以下条件，您可能希望从 Python（而不是 C++）编写自定义运算符：
您有一个 Python 函数，您希望 PyTorch 将其视为不透明的可调用对象，尤其是使用 尊重 和 。
torch.compiletorch.export您有一些 Python 绑定到 C++/CUDA 内核，并希望它们与 PyTorch 一起编写 子系统（如 或
torch.compiletorch.autograd)
将自定义 C++ 和/或 CUDA 代码与 PyTorch 集成¶
在以下情况下，您可能希望从 C++（而不是 Python）创作自定义运算符：
您有自定义 C++ 和/或 CUDA 代码。
您计划使用此代码来执行无 Python 推理。
AOTInductor
自定义操作手册¶
有关教程和本页中未涵盖的信息，请参阅 The Custom Operators Manual（我们正在努力将信息移至我们的文档站点）。我们建议您 首先阅读上面的教程之一，然后使用 Custom Operators Manual 作为参考; 它并不意味着要从头到脚阅读。
何时应创建自定义 Operator？¶
如果您的操作可表示为内置 PyTorch 运算符的组合 然后请将其编写为 Python 函数并调用它，而不是创建一个 custom 运算符。如果满足以下条件，请使用操作员注册 API 创建自定义运算符 正在调用一些 PyTorch 无法理解的库（例如自定义 C/C++ 代码， 自定义 CUDA 内核，或与 C/C++/CUDA 扩展的 Python 绑定）。
为什么要创建自定义 Operator？¶
可以通过抓取 Tensor 的数据指针来使用 C/C++/CUDA 内核 并将其传递给 pybind 的内核。但是，此方法不会与 PyTorch 子系统，如 autograd、torch.compile、vmap 等。挨次 要使用 PyTorch 子系统组合的操作，必须注册该操作 通过操作员注册 API。