PyTorch 自定義運算子¶

建立於：2024 年 6 月 18 日 | 最後更新：2025 年 1 月 06 日 | 最後驗證：2024 年 11 月 05 日

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 方面。
您有一些 C++/CUDA 核心的 Python 綁定，並希望這些綁定與 PyTorch 子系統（如 torch.compile 或 torch.autograd）組合。
您正在使用 Python（而不是像 AOTInductor 這樣僅限 C++ 的環境）。

將自定義 C++ 和/或 CUDA 程式碼與 PyTorch 整合¶

請參閱自定義 C++ 和 CUDA 運算子。

在以下情況下，您可能希望從 C++ 撰寫自定義運算子（而不是 Python）：

您有自定義 C++ 和/或 CUDA 程式碼。
您計劃將此程式碼與 AOTInductor 搭配使用以進行無 Python 推論。

自定義運算子手冊¶

如需教學和本頁未涵蓋的資訊，請參閱自定義運算子手冊（我們正在努力將資訊移至我們的文件網站）。我們建議您先閱讀上面的其中一個教學，然後將自定義運算子手冊用作參考；它並非旨在從頭到尾閱讀。

我應該何時建立自定義運算子？¶

如果您的運算可以表示為內建 PyTorch 運算子的組合，請將其編寫為 Python 函數並呼叫它，而不是建立自定義運算子。如果您要呼叫 PyTorch 無法理解的函式庫（例如，自定義 C/C++ 程式碼、自定義 CUDA 核心，或 C/C++/CUDA 擴充功能的 Python 綁定），請使用運算子註冊 API 建立自定義運算子。

我為什麼應該建立自定義運算子？¶

可以使用 C/C++/CUDA 核心，方法是取得張量的資料指標並將其傳遞給 pybind 的核心。但是，這種方法無法與 autograd、torch.compile、vmap 等 PyTorch 子系統組合。為了使運算與 PyTorch 子系統組合，必須透過運算子註冊 API 註冊它。