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 环境变量设置为非负数来控制工作人员的数量。