CUDA加速Pytorch

最新推荐文章于 2024-05-31 14:42:05 发布

Begoniaish

最新推荐文章于 2024-05-31 14:42:05 发布

阅读量1.8k

点赞数 30

文章标签： c++ 开发语言

本文链接：https://blog.csdn.net/Begoniaish/article/details/138114454

版权

利用CUDA对Pytorch模型推理加速，编写自定义算子

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

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

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

Motivation and Example

这篇笔记的其余部分将通过一个实际的示例来介绍编写和使用C++（以及CUDA）扩展的方法。如果你被追赶，或者如果你不在今天结束之前完成那个操作就会被解雇，你可以跳过这一部分，直接进入下一节的实现细节。

假设你提出了一种新型的循环单元，你发现它与现有技术相比具有更优越的性能。这种循环单元类似于LSTM，但不同之处在于它没有"遗忘门"，而是使用指数线性单元（Exponential Linear Unit，ELU）作为内部激活函数。由于这个单元永远不会忘记，我们将其称为LLTM，即Long-Long-Term-Memory单元。

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++中编写它，使用支持PyTorch后端大部分功能的ATen库，看看它是如何让我们轻松地将我们的Python代码转换过来的。然后，我们将通过将模型的部分部署到CUDA内核来进一步加速，以利用GPU提供的大规模并行性。

Writing a C++ Extension

C++扩展有两种形式：一种是使用setuptools提前构建，另一种是通过torch.utils.cpp_extension.load()进行“即时”构建。我们将从第一种方法开始，并稍后讨论后者。

Building with `setuptools`

对于“提前构建”的形式，我们通过编写一个 setup.py 脚本，使用 setuptools 来编译我们的 C++ 代码来构建我们的 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 中。

Writing the 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;
}

<torch/extension.h> 是一个一站式头文件，包含了编写 C++ 扩展所需的所有必要的 PyTorch 部分。它包括：

The ATen library, which is our primary API for tensor computation,
pybind11, which is how we create Python bindings for our C++ code,
Headers that manage the details of interaction between ATen and pybind11.

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

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

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

而不是：

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

当前关于 nvcc 错误的开放问题在此处。完整的解决方案代码示例在此处。

Forward Pass

接下来，我们可以将整个前向传播移植到 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};
}

Backward Pass

C++ 扩展 API 目前没有提供一种自动为我们生成反向函数的方式。因此，我们还必须实现 LLTM 的反向传播，该过程计算了损失相对于前向传播的每个输入的导数。最终，我们将前向和后向函数都放入一个 torch.autograd.Function 中，以创建一个良好的 Python 绑定。反向函数稍微复杂一些，所以我们不会深入挖掘代码（如果你感兴趣，可以阅读 Alex Graves 的论文，了解更多信息）：

// tanh'(z) = 1 - tanh^2(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};
}

Binding to 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++ 代码）维护扩展名，因为两者之间的不匹配可能会导致难以追踪的严重问题。

Using Your Extension

我们现在可以在 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 的编译器兼容。在实践中，这意味着你必须在 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)

Performance Comparison

既然我们能够从 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++ 实现的，因此预计会很快。尽管如此，这是一个很好的开始。

Performance on GPU Devices

关于 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 Compiling Extensions

之前我提到了构建 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 相同的信息。在后台，这将执行以下操作：

Create a temporary directory /tmp/torch_extensions/lltm,
Emit a Ninja build file into that temporary directory,
Compile your source files into a shared library,
Import this shared library as a Python module.

事实上，如果您向 cpp_extension.load() 传递 verbose=True，您将收到有关该过程的信息：

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 模块时重新加载扩展速度快，并且如果您没有更改扩展的源文件，则开销很低。

Writing a Mixed C++/CUDA extension

要真正将我们的实现提升到下一个水平，我们可以手动编写部分前向和反向传递的自定义 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 循环，你就会明白为什么这样做会快得多。

Using accessors

你可以在 CUDA 内核中看到，我们直接使用了正确类型的指针。实际上，在 CUDA 内核中直接使用高级类型不可知的张量将非常低效。

然而，这种做法以使用便捷性和可读性为代价，特别是对于高维数据而言。在我们的示例中，我们知道连续的 gates 张量有 3 个维度：

batch, size of batch_size and stride of 3*state_size
row, size of 3 and stride of state_size
index, size of state_size and stride of 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 位整数索引的打包访问器。

与 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};
}

Integrating a C++/CUDA Operation with 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
    })

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

from torch.utils.cpp_extension import load

lltm = load(name='lltm', sources=['lltm_cuda.cpp', 'lltm_cuda_kernel.cu'])

Performance Comparison

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

Forward: 149.802 us | Backward 393.458 us

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

Forward: 129.431 us | Backward 304.641 us

性能进一步提升了！

Conclusion

现在，你应该已经对 PyTorch 的 C++ 扩展机制有了良好的概述，也有了使用它们的动力。你可以在此处找到本笔记中显示的代码示例。如果有问题，请使用论坛。还请务必查看我们的 FAQ，以防遇到任何问题。

原文链接：Custom C++ and CUDA Extensions — PyTorch Tutorials 2.3.0+cu121 documentation

Begoniaish

关注

30
点赞
踩
22

收藏

觉得还不错? 一键收藏
0
评论
CUDA加速Pytorch

PyTorch提供了大量与神经网络、任意张量代数、数据整理以及其他目的相关的操作。然而，你可能仍然需要更加定制化的操作。例如，你可能想使用一种你在论文中发现的新型激活函数，或者实现你在研究中开发的操作。将这样的自定义操作集成到PyTorch中最简单的方法是通过扩展Function和Module来用Python编写它，如此处所述。这样做可以让你享受到自动微分的全部功能（免去编写导数函数的麻烦），以及Python通常的表达能力。然而，有时你的操作最好是用C++实现。
复制链接

扫一扫