Python PyTorch load_inline用法及代碼示例

本文簡要介紹python語言中 torch.utils.cpp_extension.load_inline 的用法。

用法:

torch.utils.cpp_extension.load_inline(name, cpp_sources, cuda_sources=None, functions=None, extra_cflags=None, extra_cuda_cflags=None, extra_ldflags=None, extra_include_paths=None, build_directory=None, verbose=False, with_cuda=None, is_python_module=True, with_pytorch_error_handling=True, keep_intermediates=True)

參數：

• cpp_sources-包含 C++ 源代碼的字符串或字符串列表。

• cuda_sources-包含 CUDA 源代碼的字符串或字符串列表。

• functions-要為其生成函數綁定的函數名稱列表。如果給出了字典，它應該將函數名映射到文檔字符串(否則隻是函數名)。

• with_cuda-確定是否將 CUDA 標頭和庫添加到構建中。如果設置為None(默認)，則此值將根據是否提供cuda_sources 自動確定。將其設置為 True 以強製包含 CUDA 標頭和庫。

• with_pytorch_error_handling-確定 pytorch 錯誤和警告宏是否由 pytorch 而不是 pybind 處理。為此，每個函數 foo 都通過中間函數 _safe_foo 調用。這種重定向可能會在 cpp 的晦澀情況下引起問題。當此重定向導致問題時，應將此標誌設置為 False。

從字符串源加載PyTorch C++ 擴展just-in-time (JIT)。

該函數的行為與 load() 完全相同，但將其源視為字符串而不是文件名。這些字符串存儲到構建目錄中的文件中，之後 load_inline() 的行為與 load() 相同。

請參閱the tests 以獲得使用此函數的良好示例。

源代碼可能會省略典型非內聯 C++ 擴展的兩個必需部分：必要的標頭包括，以及 (pybind11) 綁定代碼。更準確地說，傳遞給cpp_sources 的字符串首先被連接到單個.cpp 文件中。然後在此文件前麵加上 #include <torch/extension.h> 。

此外，如果提供了functions 參數，將為每個指定的函數自動生成綁定。 functions 可以是函數名列表，也可以是從函數名到文檔字符串的字典映射。如果給出一個列表，則每個函數的名稱將用作其文檔字符串。

cuda_sources 中的源被連接到一個單獨的 .cu 文件中，並在前麵加上 torch/types.h 、 cuda.h 和 cuda_runtime.h 包括。 .cpp 和 .cu 文件分別編譯，但最終鏈接到單個庫中。請注意，不會為cuda_sources 本身中的函數生成綁定。要綁定到 CUDA 內核，您必須創建一個調用它的 C++ 函數，並在 cpp_sources 之一中聲明或定義此 C++ 函數(並將其名稱包含在 functions 中)。

有關下麵省略的參數的說明，請參閱load()。

示例

>>> from torch.utils.cpp_extension import load_inline
>>> source = \'\'\'
at::Tensor sin_add(at::Tensor x, at::Tensor y) {
  return x.sin() + y.sin();
}
\'\'\'
>>> module = load_inline(name='inline_extension',
                         cpp_sources=[source],
                         functions=['sin_add'])

注意

默認情況下，Ninja 後端使用 #CPUS + 2 個工作人員來構建擴展。這可能會占用某些係統上的太多資源。可以通過將MAX_JOBS 環境變量設置為非負數來控製工作人員的數量。