自定义 C++ 和 CUDA 扩展

风尘23187

已于 2022-11-12 10:59:04 修改

阅读量1.6k

点赞数 3

CC 4.0 BY-SA版权

分类专栏：深度学习 CUDA实战 pytorch c++/cuda 文章标签： c++ python 深度学习

于 2022-11-11 10:40:58 首次发布

本文链接：https://blog.csdn.net/ygf666/article/details/127797494

来源
 官方文档

前言

PyTorch 提供了大量与神经网络、随机张量代数(arbitrary tensor algebra）、数据整合(data wrangling）以及其他目的相关的操作。但是，您仍然可能会发现自己需要更多自定义操作。例如，您可能想使用在论文中发现的新的激活函数，或者实现您在研究过程中所开发的新的运算。

在 PyTorch 中整合这样的自定义操作最简单的方法是利用 Python 编写扩展的函数(Funciton）和模型(Module），如此处所描写的那样。这让您可以充分地利用自动微分(automatic differentiation）(使你不需要自己编写派生函数）与 Python 在通常情况下的表现力。然而，在有些时候您的一些操作可以使用 C++ 以获得更佳的效果。比如，您的代码在模型当中会被十分频繁地调用，或者即便调用次数较少也会带来昂贵的开销。另一个可能的原因是您的代码依赖于一些 C 和 C++ 库，或者需要与它们交互。为了解决这种情况，PyTorch 提供了一种非常简单的编写自定义 C++ 扩展的方法。

C++ 扩展是一种我们开发的以允许用户(您）创建一些包含的资源之外的 PyTorch 运算符，例如，与 PyTorch 后端分离开来。此方法与原生的 PyTorch 操作的实现方式不同。C++ 扩展旨在为您提供大量与 PyTorch 后端集成在一起相关的样板(boilerplate），同时为基于 PyTorch 的项目提供高度的灵活性。但是，一旦将操作定义为 C++ 扩展，将其转换为原生 PyTorch 函数在很大程度上取决于您的代码组织结构，如果您决定在较早阶段进行操作，则可以解决这个问题。

1. 动机和例子

本篇文章的其余部分将逐步介绍一个编写和使用 C++(和CUDA）扩展的实际示例。如果您一直在被催促，或者在今天结束前仍未完成该扩展您就会被开除，那么可以跳过本节，直接进入下一部分的实施细节。

假设您想出了一种新型的循环单元，与现有技术相比，它具有更好的性能。该循环单元与 LSTM 相似，但不同之处在于，它没有遗忘门，并使用指数线性单元(ELU）作为其内部激活函数。由于此单元永远不会忘记，因此我们将其称为 LLTM 或长长期记忆(Long-Long-Term-Memory）单元。

由于 LLTM 和 LSTM 两者的区别过于明显，以至于我们不能通过修改 PyTorch 中的 LSTMCell 来实验我们的目标，因此我们需要创建一个自定义单元。解决这个问题的第一种也是最简单的一种 — 并且在所有情况下都是最好的一步 — 是使用 Python 在原生的 PyTorch 中实现我们所需的功能。为此，我们需要继承 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 在 NVIDIA cuDNN，Intel MKL 或 NNPACK 等库的支持下对其 CPU 和 GPU 的操作进行了高度优化的实现，因此前述的 PyTorch 代码通常足够快。但是，我们还是可以发现，在某些情况下为什么性能仍然有进一步改进的空间。最明显的原因是 PyTorch 不了解您要实现的算法。它仅知道您用于组成算法的单个操作。因此，PyTorch 必须逐个执行您的操作。由于对操作的实现(或内核）的每个单独调用(可能涉及启动CUDA内核）都具有一定的开销，因此该开销在许多函数调用中可能变得十分明显。此外，运行我们代码的 Python 解释器本身也可能会使我们的程序变慢。

一种明显的加速方法是用 C++(或CUDA）重写这部分代码并融合特定的操作组。融合是指将许多函数的实现组合到一个函数中，这可以从两个方面受益：更少的内核启动，以及在提高全局数据流可见性的情况下执行的其他优化。

让我们看看如何使用 C++ 扩展来实现 LLTM 的融合版本。我们将从使用支持 PyTorch 大部分后端功能的 ATen 库以原生 C++ 编写代码开始，然后看看它是如何让我们轻松转换 Python 代码的。然后，我们将模型的各个部分移至 CUDA 内核，以从 GPU 提供的大规模并行处理中受益，从而进一步加快处理速度。

2. 编写一个 C++ 扩展

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

2.1 使用 `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 的一个便利的包装器(wrapper），它传递正确的引用路径，并且将扩展包语言设置为 c++。等效的泛化版 setuptools 简单代码如下所示：

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

BuildExtension 执行并检查许多必需的配置步骤，并且在混合使用 C++ / CUDA 扩展的情况下管理混合编译。这就是我们目前真正需要了解的有关构建 C++ 扩展的全部信息！现在让我们看一下 lltm.cpp 中的 C++ 扩展的实现。

2.1.1 编写 c++ 操作

现在让我们开始利用 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> 是一站式(one-stop）头文件，其中包括编写 C++ 扩展所有必需的 PyTorch 扩展。这包括：

ATen 库，它是我们张量计算的主要 API，
pybind11，用于实现我们的 C++ 代码的 Python 衔接方法，
其他管理 ATen 和 pybind11 交互细节的头文件。

d_sigmoid() 的实现展示了如何使用 ATen API。PyTorch 的张量和变量接口是由 ATen 库自动生成的，因此我们可以或多或少地实现将 Python 以 1：1 的形式转换为 C++。我们用于所有计算的主要数据类型将是 torch::Tensor。它的完整 API 可以在这里查到。注意，我们可以包含 <iostream> 或任何其他 C 或 C++ 头文件 — 我们可以使用 C++11 的全部功能。

1 前向传播

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

2 后向传播

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.