PyTorch 2.2 中文官方教程（十二）（1）-阿里云开发者社区

自定义 C++和 CUDA 扩展

原文：pytorch.org/tutorials/advanced/cpp_extension.html

译者：飞龙

协议：CC BY-NC-SA 4.0

PyTorch 提供了大量与神经网络、任意张量代数、数据处理和其他目的相关的操作。然而，您可能仍然需要更定制化的操作。例如，您可能想使用在论文中找到的新型激活函数，或者实现您作为研究的一部分开发的操作。

在 PyTorch 中集成这种自定义操作的最简单方法是通过扩展Function和Module来用 Python 编写它，如此处所述。这为您提供了自动微分的全部功能（免去了编写导数函数的麻烦），以及 Python 的通常表达能力。然而，有时候您的操作最好在 C++中实现。例如，您的代码可能需要非常快，因为它在模型中被频繁调用，或者即使是少数调用也非常昂贵。另一个可能的原因是它依赖于或与其他 C 或 C++库交互。为了解决这些情况，PyTorch 提供了一种非常简单的编写自定义C++扩展的方法。

C++扩展是我们开发的一种机制，允许用户（您）创建 PyTorch 操作符定义为源外，即与 PyTorch 后端分开。这种方法与实现本机 PyTorch 操作的方式不同。C++扩展旨在为您提供高度灵活性，以便在 PyTorch 项目中节省与将操作与 PyTorch 后端集成相关的大量样板代码。然而，一旦您将操作定义为 C++扩展，将其转换为本机 PyTorch 函数在很大程度上是代码组织的问题，如果您决定向上游贡献您的操作，可以在事后处理。

动机和示例

本文的其余部分将演示如何编写和使用 C++（和 CUDA）扩展的实际示例。如果您被追赶，或者如果您不在今天结束之前完成该操作，将会被解雇，您可以跳过本节，直接前往下一节中的实现细节。

假设您想出了一种新型的循环单元，发现它具有比现有技术更优越的性能。这种循环单元类似于 LSTM，但不同之处在于它没有遗忘门，而是使用指数线性单元（ELU）作为其内部激活函数。因为这个单元永远不会忘记，我们将其称为LLTM，或长长期记忆单元。

LLTM 与普通 LSTM 不同的两种方式是显著的，以至于我们无法配置 PyTorch 的LSTMCell以满足我们的需求，因此我们必须创建一个自定义单元。这种情况下的第一种最简单的方法 - 也可能是所有情况下的一个很好的第一步 - 是在纯 PyTorch 中用 Python 实现我们想要的功能。为此，我们需要继承torch.nn.Module并实现 LLTM 的前向传播。这看起来可能是这样的：

class LLTM(torch.nn.Module):
    def __init__(self, input_features, state_size):
        super(LLTM, self).__init__()
        self.input_features = input_features
        self.state_size = state_size
        # 3 * state_size for input gate, output gate and candidate cell gate.
        # input_features + state_size because we will multiply with [input, h].
        self.weights = torch.nn.Parameter(
            torch.empty(3 * state_size, input_features + state_size))
        self.bias = torch.nn.Parameter(torch.empty(3 * state_size))
        self.reset_parameters()
    def reset_parameters(self):
        stdv = 1.0 / math.sqrt(self.state_size)
        for weight in self.parameters():
            weight.data.uniform_(-stdv, +stdv)
    def forward(self, input, state):
        old_h, old_cell = state
        X = torch.cat([old_h, input], dim=1)
        # Compute the input, output and candidate cell gates with one MM.
        gate_weights = F.linear(X, self.weights, self.bias)
        # Split the combined gate weight matrix into its components.
        gates = gate_weights.chunk(3, dim=1)
        input_gate = torch.sigmoid(gates[0])
        output_gate = torch.sigmoid(gates[1])
        # Here we use an ELU instead of the usual tanh.
        candidate_cell = F.elu(gates[2])
        # Compute the new cell state.
        new_cell = old_cell + candidate_cell * input_gate
        # Compute the new hidden state and output.
        new_h = torch.tanh(new_cell) * output_gate
        return new_h, new_cell

我们可以按预期使用：

import torch
X = torch.randn(batch_size, input_features)
h = torch.randn(batch_size, state_size)
C = torch.randn(batch_size, state_size)
rnn = LLTM(input_features, state_size)
new_h, new_C = rnn(X, (h, C))

当然，如果可能和合理的话，您应该使用这种方法来扩展 PyTorch。由于 PyTorch 高度优化了其针对 CPU 和 GPU 的操作实现，由诸如 NVIDIA cuDNN、Intel MKL 或 NNPACK 等库支持，因此像上面的 PyTorch 代码通常已经足够快。然而，我们也可以看到，在某些情况下，还有进一步提高性能的空间。最明显的原因是 PyTorch 对您正在实现的算法一无所知。它只知道您用来组成算法的各个操作。因此，PyTorch 必须逐个执行您的操作。由于对每个操作的实现（或内核）的每个单独调用，可能涉及启动 CUDA 内核，都有一定的开销，这种开销在许多函数调用中可能变得显著。此外，运行我们代码的 Python 解释器本身也可能减慢我们程序的运行速度。

因此，加快速度的一种明确方法是将部分代码重写为 C++（或 CUDA），并融合特定组的操作。融合意味着将许多函数的实现合并到一个函数中，从中获益于更少的内核启动以及我们可以通过增加数据全局流动的可见性执行的其他优化。

让我们看看如何使用 C++ 扩展来实现 LLTM 的融合版本。我们将首先用普通的 C++ 编写它，使用 ATen 库来支持 PyTorch 后端的大部分功能，并看看它是如何轻松地让我们转换我们的 Python 代码的。然后，我们将通过将模型的部分移动到 CUDA 内核来进一步加快速度，以便从 GPU 提供的大规模并行性中获益。

编写 C++ 扩展

C++ 扩展有两种类型：可以使用 setuptools “预先构建”，也可以通过 torch.utils.cpp_extension.load() “即时构建”。我们将从第一种方法开始，并稍后讨论后者。

使用 `setuptools` 构建

对于“预先构建”类型，我们通过编写一个 setup.py 脚本来构建我们的 C++ 扩展，该脚本使用 setuptools 来编译我们的 C++ 代码。对于 LLTM，它看起来就像这样简单：

from setuptools import setup, Extension
from torch.utils import cpp_extension
setup(name='lltm_cpp',
      ext_modules=[cpp_extension.CppExtension('lltm_cpp', ['lltm.cpp'])],
      cmdclass={'build_ext': cpp_extension.BuildExtension})

在这段代码中，CppExtension 是围绕 setuptools.Extension 的一个便利包装器，它传递了正确的包含路径并将扩展的语言设置为 C++。等效的原始 setuptools 代码将简单地是：

Extension(
   name='lltm_cpp',
   sources=['lltm.cpp'],
   include_dirs=cpp_extension.include_paths(),
   language='c++')

BuildExtension 执行了许多必需的配置步骤和检查，还管理了混合编译，以处理混合的 C++/CUDA 扩展。这就是我们现在需要了解有关构建 C++ 扩展的全部内容！现在让我们来看看我们的 C++ 扩展的实现，它位于 lltm.cpp 中。

编写 C++ Op

让我们开始在 C++ 中实现 LLTM！我们在反向传播中需要的一个函数是 sigmoid 的导数。这是一个足够小的代码片段，可以讨论一下在编写 C++ 扩展时可用的整体环境：

#include  <torch/extension.h>
#include  <iostream>
torch::Tensor  d_sigmoid(torch::Tensor  z)  {
  auto  s  =  torch::sigmoid(z);
  return  (1  -  s)  *  s;
}

是一个一站式头文件，包含了编写 C++ 扩展所需的所有必要 PyTorch 组件。它包括：

ATen 库是我们进行张量计算的主要 API，
pybind11 是我们为 C++ 代码创建 Python 绑定的方式，
管理 ATen 和 pybind11 之间交互细节的头文件。

d_sigmoid()的实现展示了如何使用 ATen API。PyTorch 的张量和变量接口是从 ATen 库自动生成的，因此我们可以将我们的 Python 实现几乎一对一地转换成 C++。我们所有计算的主要数据类型将是torch::Tensor。其完整 API 可以在这里查看。还要注意，我们可以包含或任何其他 C 或 C++头文件 - 我们可以充分利用 C++11 的全部功能。

请注意，CUDA-11.5 nvcc 在 Windows 上解析 torch/extension.h 时会遇到内部编译器错误。为了解决此问题，将 Python 绑定逻辑移至纯 C++文件。示例用法：

#include  <ATen/ATen.h>
at::Tensor  SigmoidAlphaBlendForwardCuda(....)

而不是：

#include  <torch/extension.h>
torch::Tensor  SigmoidAlphaBlendForwardCuda(...)

目前存在的 nvcc bug 问题请参考这里。完整的解决方案代码示例请参考这里。

前向传递

接下来我们可以将整个前向传递移植到 C++中：

#include  <vector>
std::vector<at::Tensor>  lltm_forward(
  torch::Tensor  input,
  torch::Tensor  weights,
  torch::Tensor  bias,
  torch::Tensor  old_h,
  torch::Tensor  old_cell)  {
  auto  X  =  torch::cat({old_h,  input},  /*dim=*/1);
  auto  gate_weights  =  torch::addmm(bias,  X,  weights.transpose(0,  1));
  auto  gates  =  gate_weights.chunk(3,  /*dim=*/1);
  auto  input_gate  =  torch::sigmoid(gates[0]);
  auto  output_gate  =  torch::sigmoid(gates[1]);
  auto  candidate_cell  =  torch::elu(gates[2],  /*alpha=*/1.0);
  auto  new_cell  =  old_cell  +  candidate_cell  *  input_gate;
  auto  new_h  =  torch::tanh(new_cell)  *  output_gate;
  return  {new_h,
  new_cell,
  input_gate,
  output_gate,
  candidate_cell,
  X,
  gate_weights};
}

反向传递

C++扩展 API 目前没有提供一种自动生成反向函数的方法。因此，我们还必须实现 LLTM 的反向传递，它计算损失相对于前向传递的每个输入的导数。最终，我们将前向和反向函数一起放入torch.autograd.Function中，以创建一个很好的 Python 绑定。反向函数稍微复杂一些，因此我们不会深入研究代码（如果您感兴趣，Alex Graves 的论文是一个更多信息的好读物）：

// tanh'(z) = 1 - tanh²(z)
torch::Tensor  d_tanh(torch::Tensor  z)  {
  return  1  -  z.tanh().pow(2);
}
// elu'(z) = relu'(z) + { alpha * exp(z) if (alpha * (exp(z) - 1)) < 0, else 0}
torch::Tensor  d_elu(torch::Tensor  z,  torch::Scalar  alpha  =  1.0)  {
  auto  e  =  z.exp();
  auto  mask  =  (alpha  *  (e  -  1))  <  0;
  return  (z  >  0).type_as(z)  +  mask.type_as(z)  *  (alpha  *  e);
}
std::vector<torch::Tensor>  lltm_backward(
  torch::Tensor  grad_h,
  torch::Tensor  grad_cell,
  torch::Tensor  new_cell,
  torch::Tensor  input_gate,
  torch::Tensor  output_gate,
  torch::Tensor  candidate_cell,
  torch::Tensor  X,
  torch::Tensor  gate_weights,
  torch::Tensor  weights)  {
  auto  d_output_gate  =  torch::tanh(new_cell)  *  grad_h;
  auto  d_tanh_new_cell  =  output_gate  *  grad_h;
  auto  d_new_cell  =  d_tanh(new_cell)  *  d_tanh_new_cell  +  grad_cell;
  auto  d_old_cell  =  d_new_cell;
  auto  d_candidate_cell  =  input_gate  *  d_new_cell;
  auto  d_input_gate  =  candidate_cell  *  d_new_cell;
  auto  gates  =  gate_weights.chunk(3,  /*dim=*/1);
  d_input_gate  *=  d_sigmoid(gates[0]);
  d_output_gate  *=  d_sigmoid(gates[1]);
  d_candidate_cell  *=  d_elu(gates[2]);
  auto  d_gates  =
  torch::cat({d_input_gate,  d_output_gate,  d_candidate_cell},  /*dim=*/1);
  auto  d_weights  =  d_gates.t().mm(X);
  auto  d_bias  =  d_gates.sum(/*dim=*/0,  /*keepdim=*/true);
  auto  d_X  =  d_gates.mm(weights);
  const  auto  state_size  =  grad_h.size(1);
  auto  d_old_h  =  d_X.slice(/*dim=*/1,  0,  state_size);
  auto  d_input  =  d_X.slice(/*dim=*/1,  state_size);
  return  {d_old_h,  d_input,  d_weights,  d_bias,  d_old_cell};
}

绑定到 Python

一旦您用 C++和 ATen 编写了操作，您可以使用 pybind11 以非常简单的方式将您的 C++函数或类绑定到 Python 中。关于 PyTorch C++扩展的这部分问题或问题将主要由pybind11 文档解决。

对于我们的扩展，必要的绑定代码仅涉及四行：

PYBIND11_MODULE(TORCH_EXTENSION_NAME,  m)  {
  m.def("forward",  &lltm_forward,  "LLTM forward");
  m.def("backward",  &lltm_backward,  "LLTM backward");
}

这里要注意的一点是宏TORCH_EXTENSION_NAME。torch 扩展构建将其定义为您在setup.py脚本中给出的扩展名称。在这种情况下，TORCH_EXTENSION_NAME的值将是“lltm_cpp”。这是为了避免在两个地方（构建脚本和您的 C++代码）维护扩展名，因为两者之间的不匹配可能导致难以跟踪的问题。

使用您的扩展

现在我们已经准备好在 PyTorch 中导入我们的扩展。此时，您的目录结构可能如下所示：

pytorch/
  lltm-extension/
    lltm.cpp
    setup.py

现在，运行python setup.py install来构建和安装您的扩展。这应该看起来像这样：

running install
running bdist_egg
running egg_info
creating lltm_cpp.egg-info
writing lltm_cpp.egg-info/PKG-INFO
writing dependency_links to lltm_cpp.egg-info/dependency_links.txt
writing top-level names to lltm_cpp.egg-info/top_level.txt
writing manifest file 'lltm_cpp.egg-info/SOURCES.txt'
reading manifest file 'lltm_cpp.egg-info/SOURCES.txt'
writing manifest file 'lltm_cpp.egg-info/SOURCES.txt'
installing library code to build/bdist.linux-x86_64/egg
running install_lib
running build_ext
building 'lltm_cpp' extension
creating build
creating build/temp.linux-x86_64-3.7
gcc -pthread -B ~/local/miniconda/compiler_compat -Wl,--sysroot=/ -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -I~/local/miniconda/lib/python3.7/site-packages/torch/include -I~/local/miniconda/lib/python3.7/site-packages/torch/include/torch/csrc/api/include -I~/local/miniconda/lib/python3.7/site-packages/torch/include/TH -I~/local/miniconda/lib/python3.7/site-packages/torch/include/THC -I~/local/miniconda/include/python3.7m -c lltm.cpp -o build/temp.linux-x86_64-3.7/lltm.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=lltm_cpp -D_GLIBCXX_USE_CXX11_ABI=1 -std=c++11
cc1plus: warning: command line option ‘-Wstrict-prototypes’ is valid for C/ObjC but not for C++
creating build/lib.linux-x86_64-3.7
g++ -pthread -shared -B ~/local/miniconda/compiler_compat -L~/local/miniconda/lib -Wl,-rpath=~/local/miniconda/lib -Wl,--no-as-needed -Wl,--sysroot=/ build/temp.linux-x86_64-3.7/lltm.o -o build/lib.linux-x86_64-3.7/lltm_cpp.cpython-37m-x86_64-linux-gnu.so
creating build/bdist.linux-x86_64
creating build/bdist.linux-x86_64/egg
copying build/lib.linux-x86_64-3.7/lltm_cpp.cpython-37m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/egg
creating stub loader for lltm_cpp.cpython-37m-x86_64-linux-gnu.so
byte-compiling build/bdist.linux-x86_64/egg/lltm_cpp.py to lltm_cpp.cpython-37.pyc
creating build/bdist.linux-x86_64/egg/EGG-INFO
copying lltm_cpp.egg-info/PKG-INFO -> build/bdist.linux-x86_64/egg/EGG-INFO
copying lltm_cpp.egg-info/SOURCES.txt -> build/bdist.linux-x86_64/egg/EGG-INFO
copying lltm_cpp.egg-info/dependency_links.txt -> build/bdist.linux-x86_64/egg/EGG-INFO
copying lltm_cpp.egg-info/top_level.txt -> build/bdist.linux-x86_64/egg/EGG-INFO
writing build/bdist.linux-x86_64/egg/EGG-INFO/native_libs.txt
zip_safe flag not set; analyzing archive contents...
__pycache__.lltm_cpp.cpython-37: module references __file__
creating 'dist/lltm_cpp-0.0.0-py3.7-linux-x86_64.egg' and adding 'build/bdist.linux-x86_64/egg' to it
removing 'build/bdist.linux-x86_64/egg' (and everything under it)
Processing lltm_cpp-0.0.0-py3.7-linux-x86_64.egg
removing '~/local/miniconda/lib/python3.7/site-packages/lltm_cpp-0.0.0-py3.7-linux-x86_64.egg' (and everything under it)
creating ~/local/miniconda/lib/python3.7/site-packages/lltm_cpp-0.0.0-py3.7-linux-x86_64.egg
Extracting lltm_cpp-0.0.0-py3.7-linux-x86_64.egg to ~/local/miniconda/lib/python3.7/site-packages
lltm-cpp 0.0.0 is already the active version in easy-install.pth
Installed ~/local/miniconda/lib/python3.7/site-packages/lltm_cpp-0.0.0-py3.7-linux-x86_64.egg
Processing dependencies for lltm-cpp==0.0.0
Finished processing dependencies for lltm-cpp==0.0.0

关于编译器的一点说明：由于 ABI 版本问题，用于构建 C++扩展的编译器必须与 PyTorch 构建时使用的编译器ABI 兼容。实际上，这意味着您必须在 Linux 上使用 GCC 版本 4.9 及以上。对于 Ubuntu 16.04 和其他更近期的 Linux 发行版，这应该已经是默认编译器了。在 MacOS 上，您必须使用 clang（它没有任何 ABI 版本问题）。在最坏的情况下，您可以使用您的编译器从源代码构建 PyTorch，然后使用相同的编译器构建扩展。

构建完您的扩展后，您可以在 Python 中简单地导入它，使用您在setup.py脚本中指定的名称。只需确保首先import torch，因为这将解析动态链接器必须看到的一些符号：

In [1]: import torch
In [2]: import lltm_cpp
In [3]: lltm_cpp.forward
Out[3]: <function lltm.PyCapsule.forward>

如果我们在函数或模块上调用help()，我们可以看到其签名与我们的 C++代码匹配：

In[4] help(lltm_cpp.forward)
forward(...) method of builtins.PyCapsule instance
    forward(arg0: torch::Tensor, arg1: torch::Tensor, arg2: torch::Tensor, arg3: torch::Tensor, arg4: torch::Tensor) -> List[torch::Tensor]
    LLTM forward

由于我们现在能够从 Python 调用我们的 C++函数，我们可以将它们包装在torch.autograd.Function和torch.nn.Module中，使它们成为 PyTorch 的一等公民：

import math
import torch
# Our module!
import lltm_cpp
class LLTMFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, weights, bias, old_h, old_cell):
        outputs = lltm_cpp.forward(input, weights, bias, old_h, old_cell)
        new_h, new_cell = outputs[:2]
        variables = outputs[1:] + [weights]
        ctx.save_for_backward(*variables)
        return new_h, new_cell
    @staticmethod
    def backward(ctx, grad_h, grad_cell):
        outputs = lltm_cpp.backward(
            grad_h.contiguous(), grad_cell.contiguous(), *ctx.saved_tensors)
        d_old_h, d_input, d_weights, d_bias, d_old_cell = outputs
        return d_input, d_weights, d_bias, d_old_h, d_old_cell
class LLTM(torch.nn.Module):
    def __init__(self, input_features, state_size):
        super(LLTM, self).__init__()
        self.input_features = input_features
        self.state_size = state_size
        self.weights = torch.nn.Parameter(
            torch.empty(3 * state_size, input_features + state_size))
        self.bias = torch.nn.Parameter(torch.empty(3 * state_size))
        self.reset_parameters()
    def reset_parameters(self):
        stdv = 1.0 / math.sqrt(self.state_size)
        for weight in self.parameters():
            weight.data.uniform_(-stdv, +stdv)
    def forward(self, input, state):
        return LLTMFunction.apply(input, self.weights, self.bias, *state)

性能比较

现在我们能够从 PyTorch 使用和调用我们的 C++代码，我们可以运行一个小型基准测试，看看我们从将操作重写为 C++中获得了多少性能提升。我们将运行 LLTM 的前向和反向几次，并测量持续时间：

import time
import torch
batch_size = 16
input_features = 32
state_size = 128
X = torch.randn(batch_size, input_features)
h = torch.randn(batch_size, state_size)
C = torch.randn(batch_size, state_size)
rnn = LLTM(input_features, state_size)
forward = 0
backward = 0
for _ in range(100000):
    start = time.time()
    new_h, new_C = rnn(X, (h, C))
    forward += time.time() - start
    start = time.time()
    (new_h.sum() + new_C.sum()).backward()
    backward += time.time() - start
print('Forward: {:.3f} s | Backward {:.3f} s'.format(forward, backward))

如果我们使用在本文开头纯 Python 编写的原始 LLTM 运行此代码，我们将得到以下数字（在我的机器上）：

Forward: 506.480 us | Backward 444.694 us

以及我们的新 C++版本：

Forward: 349.335 us | Backward 443.523 us

我们已经看到前向函数的显着加速（超过 30%）。对于反向函数，虽然可以看到加速，但并不是很大。我上面写的反向传播并没有特别优化，肯定可以改进。此外，PyTorch 的自动微分引擎可以自动并行化计算图，可能会使用更高效的操作流程，并且也是用 C++实现的，因此预计速度会很快。尽管如此，这是一个很好的开始。

GPU 设备上的性能

关于 PyTorch 的ATen后端的一个奇妙事实是，它抽象了您正在运行的计算设备。这意味着我们为 CPU 编写的相同代码也可以在 GPU 上运行，并且各个操作将相应地分派到针对 GPU 优化的实现。对于某些操作，如矩阵乘法（如mm或addmm），这是一个巨大的优势。让我们看看通过在 CUDA 张量上运行我们的 C++代码可以获得多少性能提升。我们不需要对实现进行任何更改，只需在 Python 中将张量放入 GPU 内存，要么在创建时添加device=cuda_device参数，要么在创建后使用.to(cuda_device)：

import torch
assert torch.cuda.is_available()
cuda_device = torch.device("cuda")  # device object representing GPU
batch_size = 16
input_features = 32
state_size = 128
# Note the device=cuda_device arguments here
X = torch.randn(batch_size, input_features, device=cuda_device)
h = torch.randn(batch_size, state_size, device=cuda_device)
C = torch.randn(batch_size, state_size, device=cuda_device)
rnn = LLTM(input_features, state_size).to(cuda_device)
forward = 0
backward = 0
for _ in range(100000):
    start = time.time()
    new_h, new_C = rnn(X, (h, C))
    torch.cuda.synchronize()
    forward += time.time() - start
    start = time.time()
    (new_h.sum() + new_C.sum()).backward()
    torch.cuda.synchronize()
    backward += time.time() - start
print('Forward: {:.3f} us | Backward {:.3f} us'.format(forward * 1e6/1e5, backward * 1e6/1e5))

再次比较我们的纯 PyTorch 代码与我们的 C++版本，现在两者都在 CUDA 设备上运行，我们再次看到性能提升。对于 Python/PyTorch：

Forward: 187.719 us | Backward 410.815 us

以及 C++/ATen：

Forward: 149.802 us | Backward 393.458 us

这是与非 CUDA 代码相比的整体加速效果很好。然而，我们可以通过编写自定义 CUDA 核心来进一步提高 C++代码的性能，我们将很快深入讨论这一点。在此之前，让我们讨论另一种构建 C++扩展的方法。

JIT 编译扩展

之前我提到构建 C++扩展有两种方法：使用setuptools或即时编译（JIT）。在介绍了前者之后，让我们详细说明后者。JIT 编译机制为您提供了一种通过调用 PyTorch API 中的一个简单函数torch.utils.cpp_extension.load()来即时编译和加载扩展的方法。对于 LLTM，这看起来就像这样简单：

from torch.utils.cpp_extension import load
lltm_cpp = load(name="lltm_cpp", sources=["lltm.cpp"])

在这里，我们为函数提供与setuptools相同的信息。在后台，这将执行以下操作：

创建一个临时目录/tmp/torch_extensions/lltm，
在临时目录中生成一个Ninja构建文件，
将您的源文件编译成共享库，
将此共享库导入为 Python 模块。

实际上，如果将verbose=True传递给cpp_extension.load()，您将了解到整个过程：

Using /tmp/torch_extensions as PyTorch extensions root...
Emitting ninja build file /tmp/torch_extensions/lltm_cpp/build.ninja...
Building extension module lltm_cpp...
Loading extension module lltm_cpp...

生成的 Python 模块将与 setuptools 生成的完全相同，但消除了必须维护单独的setup.py构建文件的要求。如果您的设置更复杂，并且确实需要setuptools的全部功能，您可以编写自己的setup.py - 但在许多情况下，这种 JIT 技术就足够了。第一次运行这行代码时，会花费一些时间，因为扩展正在后台编译。由于我们使用 Ninja 构建系统来构建您的源代码，因此重新编译是增量的，因此在第二次运行 Python 模块时重新加载扩展是快速的，如果您没有更改扩展的源文件，则开销很低。

编写混合 C++/CUDA 扩展

将我们的实现提升到下一个级别，我们可以手写部分前向和后向传递的自定义 CUDA 核心。对于 LLTM 来说，这有可能特别有效，因为有大量的逐点操作序列，可以在单个 CUDA 核心中融合并并行化。让我们看看如何编写这样一个 CUDA 核心，并使用这个扩展机制将其集成到 PyTorch 中。

编写 CUDA 扩展的一般策略是首先编写一个 C++文件，定义将从 Python 调用的函数，并使用 pybind11 将这些函数绑定到 Python。此外，这个文件还将声明在 CUDA（.cu）文件中定义的函数。然后，C++函数将进行一些检查，并最终将其调用转发到 CUDA 函数。在 CUDA 文件中，我们编写我们的实际 CUDA 核心。cpp_extension包将负责使用类似gcc的 C++编译器编译 C++源代码，使用 NVIDIA 的nvcc编译器编译 CUDA 源代码。这确保每个编译器负责编译它最擅长的文件。最终，它们将被链接成一个共享库，可以在 Python 代码中使用。

我们将从 C++文件开始，我们将称之为lltm_cuda.cpp，例如：

#include  <torch/extension.h>
#include  <vector>
// CUDA forward declarations
std::vector<torch::Tensor>  lltm_cuda_forward(
  torch::Tensor  input,
  torch::Tensor  weights,
  torch::Tensor  bias,
  torch::Tensor  old_h,
  torch::Tensor  old_cell);
std::vector<torch::Tensor>  lltm_cuda_backward(
  torch::Tensor  grad_h,
  torch::Tensor  grad_cell,
  torch::Tensor  new_cell,
  torch::Tensor  input_gate,
  torch::Tensor  output_gate,
  torch::Tensor  candidate_cell,
  torch::Tensor  X,
  torch::Tensor  gate_weights,
  torch::Tensor  weights);
// C++ interface
#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
std::vector<torch::Tensor>  lltm_forward(
  torch::Tensor  input,
  torch::Tensor  weights,
  torch::Tensor  bias,
  torch::Tensor  old_h,
  torch::Tensor  old_cell)  {
  CHECK_INPUT(input);
  CHECK_INPUT(weights);
  CHECK_INPUT(bias);
  CHECK_INPUT(old_h);
  CHECK_INPUT(old_cell);
  return  lltm_cuda_forward(input,  weights,  bias,  old_h,  old_cell);
}
std::vector<torch::Tensor>  lltm_backward(
  torch::Tensor  grad_h,
  torch::Tensor  grad_cell,
  torch::Tensor  new_cell,
  torch::Tensor  input_gate,
  torch::Tensor  output_gate,
  torch::Tensor  candidate_cell,
  torch::Tensor  X,
  torch::Tensor  gate_weights,
  torch::Tensor  weights)  {
  CHECK_INPUT(grad_h);
  CHECK_INPUT(grad_cell);
  CHECK_INPUT(input_gate);
  CHECK_INPUT(output_gate);
  CHECK_INPUT(candidate_cell);
  CHECK_INPUT(X);
  CHECK_INPUT(gate_weights);
  CHECK_INPUT(weights);
  return  lltm_cuda_backward(
  grad_h,
  grad_cell,
  new_cell,
  input_gate,
  output_gate,
  candidate_cell,
  X,
  gate_weights,
  weights);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME,  m)  {
  m.def("forward",  &lltm_forward,  "LLTM forward (CUDA)");
  m.def("backward",  &lltm_backward,  "LLTM backward (CUDA)");
}

正如您所看到的，这主要是样板代码、检查和转发到我们将在 CUDA 文件中定义的函数。我们将命名这个文件为lltm_cuda_kernel.cu（注意.cu扩展名！）。NVCC 可以合理地编译 C++11，因此我们仍然可以使用 ATen 和 C++标准库（但不能使用torch.h）。请注意，setuptools无法处理具有相同名称但不同扩展名的文件，因此如果您使用setup.py方法而不是 JIT 方法，您必须为 CUDA 文件和 C++文件分配不同的名称（对于 JIT 方法，lltm.cpp和lltm.cu将正常工作）。让我们来看一下这个文件将是什么样子：

#include  <torch/extension.h>
#include  <cuda.h>
#include  <cuda_runtime.h>
#include  <vector>
template  <typename  scalar_t>
__device__  __forceinline__  scalar_t  sigmoid(scalar_t  z)  {
  return  1.0  /  (1.0  +  exp(-z));
}

在这里，我们看到了我刚刚描述的头文件，以及我们正在使用 CUDA 特定声明，如__device__和__forceinline__，以及exp等函数。让我们继续写一些我们需要的辅助函数：

template  <typename  scalar_t>
__device__  __forceinline__  scalar_t  d_sigmoid(scalar_t  z)  {
  const  auto  s  =  sigmoid(z);
  return  (1.0  -  s)  *  s;
}
template  <typename  scalar_t>
__device__  __forceinline__  scalar_t  d_tanh(scalar_t  z)  {
  const  auto  t  =  tanh(z);
  return  1  -  (t  *  t);
}
template  <typename  scalar_t>
__device__  __forceinline__  scalar_t  elu(scalar_t  z,  scalar_t  alpha  =  1.0)  {
  return  fmax(0.0,  z)  +  fmin(0.0,  alpha  *  (exp(z)  -  1.0));
}
template  <typename  scalar_t>
__device__  __forceinline__  scalar_t  d_elu(scalar_t  z,  scalar_t  alpha  =  1.0)  {
  const  auto  e  =  exp(z);
  const  auto  d_relu  =  z  <  0.0  ?  0.0  :  1.0;
  return  d_relu  +  (((alpha  *  (e  -  1.0))  <  0.0)  ?  (alpha  *  e)  :  0.0);
}

现在实际实现一个函数，我们将再次需要两件事：一个执行我们不希望手动编写的操作并调用 CUDA 核心的函数，然后是我们想要加速的部分的实际 CUDA 核心。对于前向传递，第一个函数应该如下所示：

std::vector<torch::Tensor>  lltm_cuda_forward(
  torch::Tensor  input,
  torch::Tensor  weights,
  torch::Tensor  bias,
  torch::Tensor  old_h,
  torch::Tensor  old_cell)  {
  auto  X  =  torch::cat({old_h,  input},  /*dim=*/1);
  auto  gates  =  torch::addmm(bias,  X,  weights.transpose(0,  1));
  const  auto  batch_size  =  old_cell.size(0);
  const  auto  state_size  =  old_cell.size(1);
  auto  new_h  =  torch::zeros_like(old_cell);
  auto  new_cell  =  torch::zeros_like(old_cell);
  auto  input_gate  =  torch::zeros_like(old_cell);
  auto  output_gate  =  torch::zeros_like(old_cell);
  auto  candidate_cell  =  torch::zeros_like(old_cell);
  const  int  threads  =  1024;
  const  dim3  blocks((state_size  +  threads  -  1)  /  threads,  batch_size);
  AT_DISPATCH_FLOATING_TYPES(gates.type(),  "lltm_forward_cuda",  ([&]  {
  lltm_cuda_forward_kernel<scalar_t><<<blocks,  threads>>>(
  gates.data<scalar_t>(),
  old_cell.data<scalar_t>(),
  new_h.data<scalar_t>(),
  new_cell.data<scalar_t>(),
  input_gate.data<scalar_t>(),
  output_gate.data<scalar_t>(),
  candidate_cell.data<scalar_t>(),
  state_size);
  }));
  return  {new_h,  new_cell,  input_gate,  output_gate,  candidate_cell,  X,  gates};
}

这里的主要关注点是AT_DISPATCH_FLOATING_TYPES宏和内核启动（由<<<...>>>指示）。虽然 ATen 抽象了我们处理的张量的设备和数据类型，但在运行时，张量仍然由具体类型和具体设备的内存支持。因此，我们需要一种在运行时确定张量类型并有选择地调用具有相应正确类型签名的函数的方法。手动完成，这将（概念上）看起来像这样：

switch  (tensor.type().scalarType())  {
  case  torch::ScalarType::Double:
  return  function<double>(tensor.data<double>());
  case  torch::ScalarType::Float:
  return  function<float>(tensor.data<float>());
  ...
}

AT_DISPATCH_FLOATING_TYPES的目的是为我们处理这个分发。它接受一个类型（在我们的情况下是gates.type()），一个名称（用于错误消息）和一个 lambda 函数。在这个 lambda 函数内部，类型别名scalar_t可用，并在该上下文中定义为张量在运行时实际上是的类型。因此，如果我们有一个模板函数（我们的 CUDA 内核将是这样的），我们可以用这个scalar_t别名实例化它，正确的函数将被调用。在这种情况下，我们还想以scalar_t类型的指针形式检索张量的数据指针。如果您想要分发所有类型而不仅仅是浮点类型（Float和Double），您可以使用AT_DISPATCH_ALL_TYPES。

请注意，我们使用普通的 ATen 执行一些操作。这些操作仍将在 GPU 上运行，但使用 ATen 的默认实现。这是有道理的，因为 ATen 将使用高度优化的例程来执行矩阵乘法（例如addmm）或卷积等操作，这些操作对我们自己来说要难得多。

至于内核启动本身，我们在这里指定每个 CUDA 块将有 1024 个线程，并且整个 GPU 网格被分割为尽可能多的1 x 1024线程的块，以填充我们的矩阵，每个组件一个线程。例如，如果我们的状态大小为 2048，批处理大小为 4，我们将启动总共4 x 2 = 8个块，每个块有 1024 个线程。如果您以前从未听说过 CUDA 的“块”或“网格”，那么CUDA 的入门阅读可能会有所帮助。

实际的 CUDA 内核相当简单（如果您以前编程过 GPU 的话）：

template  <typename  scalar_t>
__global__  void  lltm_cuda_forward_kernel(
  const  scalar_t*  __restrict__  gates,
  const  scalar_t*  __restrict__  old_cell,
  scalar_t*  __restrict__  new_h,
  scalar_t*  __restrict__  new_cell,
  scalar_t*  __restrict__  input_gate,
  scalar_t*  __restrict__  output_gate,
  scalar_t*  __restrict__  candidate_cell,
  size_t  state_size)  {
  const  int  column  =  blockIdx.x  *  blockDim.x  +  threadIdx.x;
  const  int  index  =  blockIdx.y  *  state_size  +  column;
  const  int  gates_row  =  blockIdx.y  *  (state_size  *  3);
  if  (column  <  state_size)  {
  input_gate[index]  =  sigmoid(gates[gates_row  +  column]);
  output_gate[index]  =  sigmoid(gates[gates_row  +  state_size  +  column]);
  candidate_cell[index]  =  elu(gates[gates_row  +  2  *  state_size  +  column]);
  new_cell[index]  =
  old_cell[index]  +  candidate_cell[index]  *  input_gate[index];
  new_h[index]  =  tanh(new_cell[index])  *  output_gate[index];
  }
}

这里主要有趣的是，我们能够为门控矩阵中的每个单独组件完全并行计算所有这些逐点操作。如果想象要在串行中对一百万个元素进行巨大的for循环，您就会明白为什么这样会更快。

使用访问器

您可以看到在 CUDA 内核中，我们直接使用正确类型的指针进行操作。事实上，在 cuda 内核中直接使用高级类型不可知的张量将非常低效。

然而，这样做会带来易用性和可读性的代价，特别是对于高维数据。在我们的示例中，我们知道连续的gates张量有 3 个维度：

批处理，batch_size的大小和3*state_size的步幅
行，3的大小和state_size的步幅
索引，state_size的大小和步幅为1

那么我们如何在内核中访问元素gates[n][row][column]呢？事实证明，您需要步幅来使用一些简单的算术来访问您的元素。

gates.data<scalar_t>()[n*3*state_size  +  row*state_size  +  column]

除了冗长外，这个表达式需要明确知道步幅，并在其参数中传递给内核函数。您可以看到，在接受具有不同大小的多个张量的内核函数的情况下，您最终将得到一个非常长的参数列表。

对我们来说，幸运的是，ATen 提供了通过单个动态检查创建的访问器，以确保张量是指定类型和维度的。然后，访问器公开了一个 API，用于有效地访问张量元素，而无需转换为单个指针：

torch::Tensor  foo  =  torch::rand({12,  12});
// assert foo is 2-dimensional and holds floats.
auto  foo_a  =  foo.accessor<float,2>();
float  trace  =  0;
for(int  i  =  0;  i  <  foo_a.size(0);  i++)  {
  // use the accessor foo_a to get tensor data.
  trace  +=  foo_a[i][i];
}

Accessor 对象具有相对较高级的接口，具有.size()和.stride()方法以及多维索引。.accessor<>接口旨在有效地访问 cpu 张量上的数据。cuda 张量的等效物是packed_accessor64<>和packed_accessor32<>，它们产生具有 64 位或 32 位整数索引的 Packed Accessors。

Accessor 与 Packed Accessor 的根本区别在于 Packed Accessor 将大小和步幅数据复制到其结构内部，而不是指向它。这使我们能够将其传递给 CUDA 内核函数并在其中使用其接口。

我们可以设计一个函数，它接受 Packed Accessors 而不是指针。

__global__  void  lltm_cuda_forward_kernel(
  const  torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits>  gates,
  const  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  old_cell,
  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  new_h,
  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  new_cell,
  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  input_gate,
  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  output_gate,
  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  candidate_cell)

让我们分解这里使用的模板。前两个参数scalar_t和2与常规 Accessor 相同。参数torch::RestrictPtrTraits表示必须使用__restrict__关键字。还要注意，我们使用了存储大小和步幅的int32_t的PackedAccessor32变体。这很重要，因为使用 64 位变体（PackedAccessor64）可能会使内核变慢。

函数声明变为

template  <typename  scalar_t>
__global__  void  lltm_cuda_forward_kernel(
  const  torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits>  gates,
  const  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  old_cell,
  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  new_h,
  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  new_cell,
  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  input_gate,
  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  output_gate,
  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  candidate_cell)  {
  //batch index
  const  int  n  =  blockIdx.y;
  // column index
  const  int  c  =  blockIdx.x  *  blockDim.x  +  threadIdx.x;
  if  (c  <  gates.size(2)){
  input_gate[n][c]  =  sigmoid(gates[n][0][c]);
  output_gate[n][c]  =  sigmoid(gates[n][1][c]);
  candidate_cell[n][c]  =  elu(gates[n][2][c]);
  new_cell[n][c]  =
  old_cell[n][c]  +  candidate_cell[n][c]  *  input_gate[n][c];
  new_h[n][c]  =  tanh(new_cell[n][c])  *  output_gate[n][c];
  }
}

实现更加可读！然后通过在主机函数中使用.packed_accessor32<>方法创建 Packed Accessors 来调用此函数。

std::vector<torch::Tensor>  lltm_cuda_forward(
  torch::Tensor  input,
  torch::Tensor  weights,
  torch::Tensor  bias,
  torch::Tensor  old_h,
  torch::Tensor  old_cell)  {
  auto  X  =  torch::cat({old_h,  input},  /*dim=*/1);
  auto  gate_weights  =  torch::addmm(bias,  X,  weights.transpose(0,  1));
  const  auto  batch_size  =  old_cell.size(0);
  const  auto  state_size  =  old_cell.size(1);
  auto  gates  =  gate_weights.reshape({batch_size,  3,  state_size});
  auto  new_h  =  torch::zeros_like(old_cell);
  auto  new_cell  =  torch::zeros_like(old_cell);
  auto  input_gate  =  torch::zeros_like(old_cell);
  auto  output_gate  =  torch::zeros_like(old_cell);
  auto  candidate_cell  =  torch::zeros_like(old_cell);
  const  int  threads  =  1024;
  const  dim3  blocks((state_size  +  threads  -  1)  /  threads,  batch_size);
  AT_DISPATCH_FLOATING_TYPES(gates.type(),  "lltm_forward_cuda",  ([&]  {
  lltm_cuda_forward_kernel<scalar_t><<<blocks,  threads>>>(
  gates.packed_accessor32<scalar_t,3,torch::RestrictPtrTraits>(),
  old_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
  new_h.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
  new_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
  input_gate.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
  output_gate.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
  candidate_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>());
  }));
  return  {new_h,  new_cell,  input_gate,  output_gate,  candidate_cell,  X,  gates};
}

反向传播遵循了大致相同的模式，我不会进一步详细说明：

template  <typename  scalar_t>
__global__  void  lltm_cuda_backward_kernel(
  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  d_old_cell,
  torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits>  d_gates,
  const  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  grad_h,
  const  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  grad_cell,
  const  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  new_cell,
  const  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  input_gate,
  const  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  output_gate,
  const  torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>  candidate_cell,
  const  torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits>  gate_weights)  {
  //batch index
  const  int  n  =  blockIdx.y;
  // column index
  const  int  c  =  blockIdx.x  *  blockDim.x  +  threadIdx.x;
  if  (c  <  d_gates.size(2)){
  const  auto  d_output_gate  =  tanh(new_cell[n][c])  *  grad_h[n][c];
  const  auto  d_tanh_new_cell  =  output_gate[n][c]  *  grad_h[n][c];
  const  auto  d_new_cell  =
  d_tanh(new_cell[n][c])  *  d_tanh_new_cell  +  grad_cell[n][c];
  d_old_cell[n][c]  =  d_new_cell;
  const  auto  d_candidate_cell  =  input_gate[n][c]  *  d_new_cell;
  const  auto  d_input_gate  =  candidate_cell[n][c]  *  d_new_cell;
  d_gates[n][0][c]  =
  d_input_gate  *  d_sigmoid(gate_weights[n][0][c]);
  d_gates[n][1][c]  =
  d_output_gate  *  d_sigmoid(gate_weights[n][1][c]);
  d_gates[n][2][c]  =
  d_candidate_cell  *  d_elu(gate_weights[n][2][c]);
  }
}
std::vector<torch::Tensor>  lltm_cuda_backward(
  torch::Tensor  grad_h,
  torch::Tensor  grad_cell,
  torch::Tensor  new_cell,
  torch::Tensor  input_gate,
  torch::Tensor  output_gate,
  torch::Tensor  candidate_cell,
  torch::Tensor  X,
  torch::Tensor  gates,
  torch::Tensor  weights)  {
  auto  d_old_cell  =  torch::zeros_like(new_cell);
  auto  d_gates  =  torch::zeros_like(gates);
  const  auto  batch_size  =  new_cell.size(0);
  const  auto  state_size  =  new_cell.size(1);
  const  int  threads  =  1024;
  const  dim3  blocks((state_size  +  threads  -  1)  /  threads,  batch_size);
  AT_DISPATCH_FLOATING_TYPES(X.type(),  "lltm_backward_cuda",  ([&]  {
  lltm_cuda_backward_kernel<scalar_t><<<blocks,  threads>>>(
  d_old_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
  d_gates.packed_accessor32<scalar_t,3,torch::RestrictPtrTraits>(),
  grad_h.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
  grad_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
  new_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
  input_gate.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
  output_gate.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
  candidate_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
  gates.packed_accessor32<scalar_t,3,torch::RestrictPtrTraits>());
  }));
  auto  d_gate_weights  =  d_gates.reshape({batch_size,  3*state_size});
  auto  d_weights  =  d_gate_weights.t().mm(X);
  auto  d_bias  =  d_gate_weights.sum(/*dim=*/0,  /*keepdim=*/true);
  auto  d_X  =  d_gate_weights.mm(weights);
  auto  d_old_h  =  d_X.slice(/*dim=*/1,  0,  state_size);
  auto  d_input  =  d_X.slice(/*dim=*/1,  state_size);
  return  {d_old_h,  d_input,  d_weights,  d_bias,  d_old_cell,  d_gates};
}

将 C++/CUDA 操作集成到 PyTorch 中

再次非常简单地将我们的 CUDA 启用的操作集成到 PyTorch 中。如果您想编写一个setup.py脚本，它可能如下所示：

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
setup(
    name='lltm',
    ext_modules=[
        CUDAExtension('lltm_cuda', [
            'lltm_cuda.cpp',
            'lltm_cuda_kernel.cu',
        ])
    ],
    cmdclass={
        'build_ext': BuildExtension
    })

现在我们使用CUDAExtension()而不是CppExtension()。我们只需指定.cu文件以及.cpp文件 - 库会为您处理所有这些麻烦。JIT 机制甚至更简单：

from torch.utils.cpp_extension import load
lltm = load(name='lltm', sources=['lltm_cuda.cpp', 'lltm_cuda_kernel.cu'])

性能比较

我们希望通过将代码的逐点操作并行化和融合到 CUDA 中，可以提高 LLTM 的性能。让我们看看这是否成立。我们可以运行我之前列出的代码来运行基准测试。我们之前最快的版本是基于 CUDA 的 C++代码：

Forward: 149.802 us | Backward 393.458 us

现在使用我们自定义的 CUDA 内核：

Forward: 129.431 us | Backward 304.641 us

更多性能提升！

结论

现在，您应该已经掌握了 PyTorch 的 C++扩展机制的概述以及使用它们的动机。您可以在此笔记中找到显示的代码示例这里。如果您有问题，请使用论坛。还要确保查看我们的FAQ，以防遇到任何问题。

使用自定义 C++运算符扩展 TorchScript

原文：pytorch.org/tutorials/advanced/torch_script_custom_ops.html

译者：飞龙

协议：CC BY-NC-SA 4.0

PyTorch 1.0 发布引入了一个称为TorchScript的新编程模型到 PyTorch 中。TorchScript 是 Python 编程语言的一个子集，可以被 TorchScript 编译器解析、编译和优化。此外，编译后的 TorchScript 模型可以选择被序列化为磁盘文件格式，然后可以在纯 C++（以及 Python）中加载和运行进行推断。

TorchScript 支持torch包提供的大量操作的子集，允许您纯粹将许多种复杂模型表达为 PyTorch“标准库”中的一系列张量操作。然而，可能会有时候您需要扩展 TorchScript 以使用自定义的 C++或 CUDA 函数。虽然我们建议只有在您的想法无法（足够高效地）表达为简单的 Python 函数时才使用此选项，但我们提供了一个非常友好和简单的接口来使用ATen，PyTorch 的高性能 C++张量库来定义自定义的 C++和 CUDA 核心。一旦绑定到 TorchScript 中，您可以将这些自定义核心（或“ops”）嵌入到您的 TorchScript 模型中，并在 Python 中执行它们，也可以直接在 C++中执行它们的序列化形式。

以下段落给出了一个编写 TorchScript 自定义操作的示例，以调用OpenCV，这是一个用 C++编写的计算机视觉库。我们将讨论如何在 C++中处理张量，如何高效地将它们转换为第三方张量格式（在本例中为 OpenCV Mat），如何在 TorchScript 运行时注册您的运算符，最后如何编译运算符并在 Python 和 C++中使用它。

在 C++中实现自定义运算符

在本教程中，我们将暴露warpPerspective函数，该函数将 OpenCV 中的透视变换应用于图像，将其作为自定义运算符从 OpenCV 到 TorchScript。第一步是在 C++中编写我们自定义运算符的实现。让我们将此实现的文件命名为op.cpp，并使其如下所示：

torch::Tensor  warp_perspective(torch::Tensor  image,  torch::Tensor  warp)  {
  // BEGIN image_mat
  cv::Mat  image_mat(/*rows=*/image.size(0),
  /*cols=*/image.size(1),
  /*type=*/CV_32FC1,
  /*data=*/image.data_ptr<float>());
  // END image_mat
  // BEGIN warp_mat
  cv::Mat  warp_mat(/*rows=*/warp.size(0),
  /*cols=*/warp.size(1),
  /*type=*/CV_32FC1,
  /*data=*/warp.data_ptr<float>());
  // END warp_mat
  // BEGIN output_mat
  cv::Mat  output_mat;
  cv::warpPerspective(image_mat,  output_mat,  warp_mat,  /*dsize=*/{8,  8});
  // END output_mat
  // BEGIN output_tensor
  torch::Tensor  output  =  torch::from_blob(output_mat.ptr<float>(),  /*sizes=*/{8,  8});
  return  output.clone();
  // END output_tensor
}

这个运算符的代码非常简短。在文件顶部，我们包含了 OpenCV 头文件opencv2/opencv.hpp，以及torch/script.h头文件，后者从 PyTorch 的 C++ API 中暴露了我们编写自定义 TorchScript 运算符所需的所有必要内容。我们的函数warp_perspective接受两个参数：一个输入image和我们希望应用于图像的warp变换矩阵。这些输入的类型是torch::Tensor，PyTorch 在 C++中的张量类型（也是 Python 中所有张量的基础类型）。我们的warp_perspective函数的返回类型也将是torch::Tensor。

提示

有关 ATen 的更多信息，请参阅此说明，该说明提供了Tensor类给 PyTorch。此外，此教程描述了如何在 C++中分配和初始化新的张量对象（对于此运算符不是必需的）。

注意

TorchScript 编译器了解固定数量的类型。只有这些类型可以用作自定义运算符的参数。目前这些类型是：torch::Tensor、torch::Scalar、double、int64_t和这些类型的std::vector。请注意只有double而不是float，只有int64_t而不是其他整数类型如int、short或long被支持。

在我们的函数内部，我们需要做的第一件事是将我们的 PyTorch 张量转换为 OpenCV 矩阵，因为 OpenCV 的warpPerspective期望cv::Mat对象作为输入。幸运的是，有一种方法可以不复制任何数据来做到这一点。在前几行中，

cv::Mat  image_mat(/*rows=*/image.size(0),
  /*cols=*/image.size(1),
  /*type=*/CV_32FC1,
  /*data=*/image.data_ptr<float>());

我们正在调用 OpenCV Mat类的此构造函数来将我们的张量转换为Mat对象。我们传递原始image张量的行数和列数，数据类型（在本例中我们将其固定为float32），最后是底层数据的原始指针 - 一个float*。Mat类的这个构造函数的特殊之处在于它不会复制输入数据。相反，它将简单地引用这个内存，用于对Mat执行的所有操作。如果在image_mat上执行了原位操作，这将反映在原始image张量中（反之亦然）。这使我们能够使用库的本机矩阵类型调用后续的 OpenCV 例程，即使我们实际上是在 PyTorch 张量中存储数据。我们重复这个过程将warp PyTorch 张量转换为warp_mat OpenCV 矩阵：

cv::Mat  warp_mat(/*rows=*/warp.size(0),
  /*cols=*/warp.size(1),
  /*type=*/CV_32FC1,
  /*data=*/warp.data_ptr<float>());

接下来，我们准备调用我们在 TorchScript 中急切想要使用的 OpenCV 函数：warpPerspective。为此，我们将image_mat和warp_mat矩阵以及一个名为output_mat的空输出矩阵传递给 OpenCV 函数。我们还指定了我们希望输出矩阵（图像）的大小dsize。在本例中，它被硬编码为8 x 8：

cv::Mat  output_mat;
  cv::warpPerspective(image_mat,  output_mat,  warp_mat,  /*dsize=*/{8,  8});

我们自定义运算符实现的最后一步是将output_mat转换回 PyTorch 张量，以便我们可以在 PyTorch 中进一步使用它。这与我们之前转换的过程非常相似。在这种情况下，PyTorch 提供了一个torch::from_blob方法。在这种情况下，blob意味着我们希望将其解释为 PyTorch 张量的一些不透明的、扁平的内存指针。调用torch::from_blob看起来像这样：

torch::Tensor  output  =  torch::from_blob(output_mat.ptr<float>(),  /*sizes=*/{8,  8});
  return  output.clone();

我们使用 OpenCV 的Mat类上的.ptr()方法来获取底层数据的原始指针（就像之前 PyTorch 张量的.data_ptr()一样）。我们还指定了张量的输出形状，我们将其硬编码为8 x 8。torch::from_blob的输出是一个指向 OpenCV 矩阵所拥有内存的torch::Tensor。

在从我们的运算符实现中返回这个张量之前，我们必须在张量上调用.clone()来执行底层数据的内存复制。这样做的原因是torch::from_blob返回一个不拥有数据的张量。此时，数据仍然由 OpenCV 矩阵拥有。然而，这个 OpenCV 矩阵将在函数结束时超出范围并被释放。如果我们原样返回output张量，那么在函数外部使用时它将指向无效的内存。调用.clone()返回一个新的张量，其中包含原始数据的副本，新张量自己拥有。因此，可以安全地返回到外部世界。

流畅的 Python 第二版（GPT 重译）（八）（2）https://developer.aliyun.com/article/1482558

PyTorch 2.2 中文官方教程（十二）（1）

自定义 C++和 CUDA 扩展

动机和示例

编写 C++ 扩展

使用 `setuptools` 构建

编写 C++ Op

前向传递

反向传递

绑定到 Python

使用您的扩展

性能比较

GPU 设备上的性能

JIT 编译扩展

编写混合 C++/CUDA 扩展

使用访问器

将 C++/CUDA 操作集成到 PyTorch 中

性能比较

结论

使用自定义 C++运算符扩展 TorchScript

在 C++中实现自定义运算符

热门文章

最新文章

相关课程

相关电子书

相关实验场景

推荐镜像

热门

活动广场

任务中心

开发者评测

高校计划

乘风者计划

训练营

阿里云MVP

话题

直播

下载

镜像站

技术资料

插件

PyTorch 2.2 中文官方教程（十二）（1）

自定义 C++和 CUDA 扩展

动机和示例

编写 C++ 扩展

使用 setuptools 构建

编写 C++ Op

前向传递

反向传递

绑定到 Python

使用您的扩展

性能比较

GPU 设备上的性能

JIT 编译扩展

编写混合 C++/CUDA 扩展

使用访问器

将 C++/CUDA 操作集成到 PyTorch 中

性能比较

结论

使用自定义 C++运算符扩展 TorchScript

在 C++中实现自定义运算符

热门文章

最新文章

相关课程

相关电子书

相关实验场景

推荐镜像

使用 `setuptools` 构建