PyTorch 自定義 C ++和 CUDA 擴展

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

作者： Peter Goldsborough

PyTorch 提供了與神經(jīng)網(wǎng)絡(luò)，任意張量代數(shù)，數(shù)據(jù)整理和其他目的有關(guān)的大量操作。 但是，您仍然可能會發(fā)現(xiàn)自己需要更多的自定義操作。 例如，您可能想使用論文中發(fā)現(xiàn)的新穎的激活功能，或者實現(xiàn)您在研究過程中開發(fā)的操作。

在 PyTorch 中集成這樣的自定義操作的最簡單方法是通過擴展此處概述的Function和Module在 Python 中編寫它。 這為您提供了自動區(qū)分的全部功能(使您不必編寫派生函數(shù)）以及 Python 的通常表達能力。 但是，有時您的操作可以用 C ++更好地實現(xiàn)。 例如，您的代碼可能需要確實快速，因為在模型中它經(jīng)常被調(diào)用，或者即使很少調(diào)用也很昂貴。 另一個合理的原因是它依賴于其他 C 或 C ++庫或與之交互。 為了解決這種情況，PyTorch 提供了一種非常簡單的方式來編寫自定義 C ++擴展。

C ++擴展是我們開發(fā)的一種機制，允許用戶(您）創(chuàng)建源外定義的 PyTorch 運算符，即，即與 PyTorch 后端分開。 該方法與與不同于本機 PyTorch 操作的實現(xiàn)方式。 C ++擴展旨在為您節(jié)省大量與將操作與 PyTorch 后端集成在一起相關(guān)的樣板，同時為基于 PyTorch 的項目提供高度的靈活性。 但是，一旦您將操作定義為 C ++擴展，將其轉(zhuǎn)換為本地 PyTorch 函數(shù)在很大程度上取決于代碼組織，如果您決定在上游進行操作，則可以解決此問題。

動機與榜樣

本說明的其余部分將逐步介紹編寫和使用 C ++(和 CUDA）擴展的實際示例。 如果您被追捕，或者在一天結(jié)束前仍未完成該操作，就會有人開除您，則可以跳過本節(jié)，直接進入下一部分的實施細節(jié)。

假設(shè)您想出了一種新型的循環(huán)裝置，發(fā)現(xiàn)與現(xiàn)有技術(shù)相比，它具有更好的性能。 該循環(huán)單元類似于 LSTM，但不同之處在于它缺少遺忘門，并使用指數(shù)線性單元(ELU）作為其內(nèi)部激活功能。 由于此設(shè)備永遠不會忘記，因此我們將其稱為 LLTM 或長期內(nèi)存單元。

LLTM 與普通 LSTM 的兩種區(qū)別非常重要，以至于我們無法為自己的目的配置 PyTorch 的LSTMCell，因此我們必須創(chuàng)建一個自定義單元。 這樣做的第一個也是最簡單的方法，并且在所有情況下都可能是一個好的第一步，是使用 Python 在純 PyTorch 中實現(xiàn)我們所需的功能。 為此，我們需要繼承torch.nn.Module，并實現(xiàn) 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

然后我們可以按預(yù)期使用：

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))

自然，如果可能的話，您應(yīng)該使用這種方法擴展 PyTorch。 由于 PyTorch 對 CPU 和 GPU 的操作進行了高度優(yōu)化的實現(xiàn)，并由 NVIDIA cuDNN ， Intel MKL 或  NNPACK 等庫提供支持 ，上面的 PyTorch 代碼通常會足夠快。 但是，我們還可以看到為什么在某些情況下還有進一步改進性能的空間。 最明顯的原因是 PyTorch 不了解您要實現(xiàn)的算法。 它僅知道您用于組成算法的單個操作。 因此，PyTorch 必須一個接一個地執(zhí)行您的操作。 由于對操作的實現(xiàn)(或內(nèi)核）的每個單獨調(diào)用(可能涉及 CUDA 內(nèi)核的啟動）都具有一定的開銷，因此該開銷在許多函數(shù)調(diào)用中可能變得很重要。 此外，運行我們的代碼的 Python 解釋器本身可能會使我們的程序變慢。

因此，一種確定的加速方法是用 C ++(或 CUDA）和熔斷特定操作組來重寫零件。 融合意味著將許多功能的實現(xiàn)組合到一個功能中，這可以從更少的內(nèi)核啟動以及我們可以通過提高全局數(shù)據(jù)流可見性而執(zhí)行的其他優(yōu)化中獲利。

讓我們看看如何使用 C ++擴展來實現(xiàn) LLTM 的融合版本。 首先，我們使用 ATen 庫以普通的 C ++語言編寫代碼，該庫為 PyTorch 的許多后端提供了強大的支持，并了解它使我們輕松轉(zhuǎn)換 Python 代碼的方式。 然后，我們將模型的某些部分移至 CUDA 內(nèi)核，以從 GPU 提供的大量并行處理中受益，從而進一步加快處理速度。

編寫 C ++擴展

C ++擴展有兩種形式：它們可以使用setuptools提前構(gòu)建，也可以通過torch.utils.cpp_extension.load()適時構(gòu)建。 我們將從第一種方法開始，稍后再討論后者。

使用setuptools構(gòu)建

為了“提前”，我們通過編寫一個setup.py腳本來構(gòu)建 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的便利包裝，它傳遞正確的包含路徑并將擴展語言設(shè)置為 C ++。 等效的原始setuptools代碼將是：

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

BuildExtension執(zhí)行許多必需的配置步驟，并檢查和管理混合 C ++ / CUDA 擴展的混合編譯。 這就是我們現(xiàn)在真正需要了解的有關(guān)構(gòu)建 C ++擴展的全部信息！ 現(xiàn)在讓我們看一下lltm.cpp中 C ++擴展的實現(xiàn)。

編寫 C ++ Op

讓我們開始以 C ++實現(xiàn) LLTM！ 我們需要向后傳遞的一項功能是 S 形導(dǎo)數(shù)。 這是一小段代碼，用于討論編寫 C ++擴展時可供我們使用的總體環(huán)境：

#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 位。 這包括：

• ATen 庫，這是我們用于張量計算的主要 API，
• pybind11 ，這是我們?yōu)?C ++代碼創(chuàng)建 Python 綁定的方式，
• 標頭，用于管理 ATen 與 pybind11 之間的交互的詳細信息。

d_sigmoid()的實現(xiàn)顯示了如何使用 ATen API。 PyTorch 的張量和變量接口是從 ATen 庫自動生成的，因此我們可以將 Python 實現(xiàn) 1：1 或多或少地轉(zhuǎn)換為 C ++。 我們用于所有計算的主要數(shù)據(jù)類型將為torch::Tensor。 可以在中檢查其完整的 API。 還要注意，我們可以包括<iostream>或任何其他 C 或 C ++頭文件 –我們擁有 C ++ 11 的全部功能。

前進通行證

接下來，我們可以將整個正向傳遞到 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 當前不提供為我們自動生成向后函數(shù)的方法。 因此，我們還必須實現(xiàn) LLTM 的后向傳遞，它計算相對于前向傳遞的每個輸入的損耗導(dǎo)數(shù)。 最終，我們將前進和后退功能放入torch.autograd.Function中，以創(chuàng)建一個不錯的 Python 綁定。 向后函數(shù)的功能稍微復(fù)雜一些，因此我們將不深入研究代碼(如果您有興趣，請閱讀 Alex Graves 的論文，以獲取有關(guān)此方面的更多信息）：

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

綁定到 Python

一旦用 C ++和 ATen 編寫了操作，就可以使用 pybind11 以非常簡單的方式將 C ++函數(shù)或類綁定到 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擴展程序構(gòu)建會將其定義為您在setup.py腳本中為擴展程序指定的名稱。 在這種情況下，TORCH_EXTENSION_NAME的值為“ lltm”。 這是為了避免必須在兩個位置(構(gòu)建腳本和 C ++代碼）維護擴展名，因為兩者之間的不匹配會導(dǎo)致令人討厭且難以跟蹤的問題。

使用擴展

現(xiàn)在，我們準備將擴展名導(dǎo)入 PyTorch 中。 此時，目錄結(jié)構(gòu)可能如下所示：

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

現(xiàn)在，運行python setup.py install來構(gòu)建和安裝擴展程序。 看起來應(yīng)該像這樣：

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

關(guān)于編譯器的小注釋：由于 ABI 版本問題，用于構(gòu)建 C ++擴展的編譯器必須為，并且 PyTorch 編譯器是與 ABI 兼容的。 實際上，這意味著您必須在 Linux 上使用 GCC 4.9 及更高版本。 對于 Ubuntu 16.04 和其他較新的 Linux 發(fā)行版，這應(yīng)該已經(jīng)是默認編譯器。 在 MacOS 上，您必須使用 clang(它沒有任何 ABI 版本控制問題）。 在最壞的情況下，您可以使用編譯器從源代碼構(gòu)建 PyTorch，然后使用相同的編譯器構(gòu)建擴展。

擴展程序構(gòu)建完成后，您可以使用在setup.py腳本中指定的名稱，簡單地將其導(dǎo)入 Python。 只需確保先import torch，因為這將解決動態(tài)鏈接器必須看到的一些符號：

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

如果在函數(shù)或模塊上調(diào)用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

由于我們現(xiàn)在可以從 Python 調(diào)用 C ++函數(shù)，因此可以將它們包裝為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_variables)
        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)

性能比較

既然我們已經(jīng)能夠使用和調(diào)用 PyTorch 的 C ++代碼，我們就可以運行一個小型基準測試，以查看通過用 C ++重寫 op 獲得的性能。 我們將向前和向后運行 LLTM 幾次，并測量持續(xù)時間：

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} us | Backward {:.3f} us'.format(forward * 1e6/1e5, backward * 1e6/1e5))

如果我們使用本文開頭用純 Python 編寫的原始 LLTM 來運行此代碼，則會得到以下數(shù)字(在我的機器上）：

Forward: 506.480 us | Backward 444.694 us

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

Forward: 349.335 us | Backward 443.523 us

我們已經(jīng)可以看到前進功能的顯著提速(超過 30％）。 對于后退功能，可以看到加速，盡管不是主要的。 我在上面編寫的后向通行證沒有特別優(yōu)化，并且肯定可以改進。 而且，PyTorch 的自動微分引擎可以自動并行化計算圖，可以整體上使用更高效的操作流程，并且也可以用 C ++實現(xiàn)，因此有望實現(xiàn)更快的速度。 盡管如此，這是一個良好的開始。

GPU 設(shè)備上的性能

關(guān)于 PyTorch 的 ATen 后端的一個奇妙事實是，它抽象了您正在運行的計算設(shè)備。 這意味著我們?yōu)?CPU 編寫的相同代碼可以也可以在 GPU 上運行，并且各個操作將相應(yīng)地分派到 GPU 優(yōu)化的實現(xiàn)。 對于某些運算，例如矩陣乘法(例如mm或addmm），這是一個很大的勝利。 讓我們看一下使用 CUDA 張量運行 C ++代碼所獲得的性能。 無需更改實現(xiàn)，只需將張量從 Python 放到 GPU 內(nèi)存中，在創(chuàng)建時添加device=cuda_device參數(shù)，或者在創(chuàng)建后使用.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 ++版本(現(xiàn)在都在 CUDA 設(shè)備上運行）進行比較，我們再次看到了性能提升。 對于 Python / PyTorch：

Forward: 187.719 us | Backward 410.815 us

和 C ++ / ATen：

Forward: 149.802 us | Backward 393.458 us

與非 CUDA 代碼相比，這可以大大提高整體速度。 但是，通過編寫自定義 CUDA 內(nèi)核，我們可以從 C ++代碼中獲得更多性能，我們將很快深入其中。 在此之前，讓我們討論構(gòu)建 C ++擴展的另一種方法。

JIT 編譯擴展

之前，我提到過有兩種構(gòu)建 C ++擴展的方法：使用setuptools或即時(JIT）。 在介紹了前者之后，讓我們詳細介紹后者。 JIT 編譯機制通過調(diào)用 PyTorch API 中稱為torch.utils.cpp_extension.load()的簡單函數(shù)，為您動態(tài)編譯和加載擴展程序。 對于 LLTM，這看起來像這樣簡單：

from torch.utils.cpp_extension import load


lltm_cpp = load(name="lltm_cpp", sources=["lltm.cpp"])

在此，我們?yōu)楹瘮?shù)提供與setuptools相同的信息。 在后臺，這將執(zhí)行以下操作：

1. 創(chuàng)建一個臨時目錄/tmp/torch_extensions/lltm，
2. 將 Ninja 構(gòu)建文件發(fā)送到該臨時目錄中，
3. 將您的源文件編譯到共享庫中，
4. 將此共享庫導(dǎo)入為 Python 模塊。

實際上，如果將verbose=True傳遞給cpp_extension.load()，則會通知您有關(guān)過程：

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構(gòu)建文件的要求。 如果您的設(shè)置更加復(fù)雜，并且確實需要setuptools的全部功能，則可以編寫自己的setup.py –但是在許多情況下，這種 JIT 技術(shù)就可以了。 第一次運行此行時，將需要一些時間，因為擴展程序是在后臺編譯的。 由于我們使用 Ninja 構(gòu)建系統(tǒng)來構(gòu)建您的源代碼，因此重新編譯是增量的，因此在您第二次運行 Python 模塊時重新加載擴展程序非?？旖荩胰绻桓臄U展程序的源文件，則開銷很低。

編寫混合的 C ++ / CUDA 擴展

為了將實現(xiàn)真正提升到一個新的水平，我們可以使用自定義 CUDA 內(nèi)核來手寫前進和后退傳遞的部分內(nèi)容。 對于 LLTM，這具有特別有效的前景，因為有大量按順序進行的逐點操作，這些操作都可以在單個 CUDA 內(nèi)核中融合和并行化。 讓我們看看如何編寫這種 CUDA 內(nèi)核，并使用此擴展機制將其與 PyTorch 集成。

編寫 CUDA 擴展的一般策略是首先編寫一個 C ++文件，該文件定義將從 Python 調(diào)用的函數(shù)，然后使用 pybind11 將這些函數(shù)綁定到 Python。 此外，此文件還將聲明在 CUDA(.cu）文件中定義的函數(shù)。 然后，C ++函數(shù)將進行一些檢查，并最終將其調(diào)用轉(zhuǎn)發(fā)給 CUDA 函數(shù)。 在 CUDA 文件中，我們編寫了實際的 CUDA 內(nèi)核。 然后cpp_extension包將負責(zé)使用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.type().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)");
}

如您所見，它主要是樣板文件，檢查并轉(zhuǎn)發(fā)到我們將在 CUDA 文件中定義的功能。 我們將此文件命名為lltm_cuda_kernel.cu(請注意.cu/擴展名！）。 NVCC 可以合理地編譯 C ++ 11，因此我們?nèi)匀豢梢允褂?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__）以及函數(shù)(例如exp）的事實。 讓我們繼續(xù)一些我們需要的輔助功??能：

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

現(xiàn)在要真正實現(xiàn)一個函數(shù)，我們將再次需要兩件事：一個函數(shù)執(zhí)行我們不希望手工編寫并調(diào)用 CUDA 內(nèi)核的操作，然后是要加速的部分的實際 CUDA 內(nèi)核。 。 對于前向傳遞，第一個函數(shù)應(yīng)如下所示：

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

這里的主要關(guān)注點是AT_DISPATCH_FLOATING_TYPES宏和內(nèi)核啟動(由<<<...>>>指示）。 盡管 ATen 提取了我們處理過的張量的設(shè)備和數(shù)據(jù)類型，但張量在運行時仍將由具體設(shè)備上具體類型的內(nèi)存支持。 因此，我們需要一種在運行時確定張量是什么類型，然后有選擇地調(diào)用具有相應(yīng)正確類型簽名的函數(shù)的方法。 手動完成后，(在概念上）將如下所示：

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的目的是為我們處理此調(diào)度。 它需要一個類型(在我們的示例中為gates.type()），一個名稱(用于錯誤消息）和一個 lambda 函數(shù)。 在此 lambda 函數(shù)中，類型別名scalar_t可用，并且定義為該上下文中張量實際上在運行時的類型。 這樣，如果我們有一個模板函數(shù)(CUDA 內(nèi)核將使用它），則可以使用此scalar_t別名實例化它，然后將調(diào)用正確的函數(shù)。 在這種情況下，我們還希望檢索張量的數(shù)據(jù)指針作為scalar_t類型的指針。 如果您想分派所有類型而不僅僅是浮點類型(Float和Double），則可以使用AT_DISPATCH_ALL_TYPES。

請注意，我們使用普通的 ATen 執(zhí)行一些操作。 這些操作仍將在 GPU 上運行，但使用 ATen 的默認實現(xiàn)。 這是有道理的，因為 ATen 會針對矩陣乘法(例如addmm）或卷積使用高度優(yōu)化的例程，而這將很難實現(xiàn)和改善。

至于內(nèi)核啟動本身，我們在這里指定每個 CUDA 塊將具有 1024 個線程，并且將整個 GPU 網(wǎng)格分為所需的1 x 1024線程塊，以便用每個組件一個線程填充矩陣。 例如，如果我們的狀態(tài)大小為 2048，批處理大小為 4，則我們將以每 1024 個線程總共啟動4 x 2 = 8塊。 如果您以前從未聽說過 CUDA 的“障礙”或“網(wǎng)格”，那么簡介 CUDA 可能會有所幫助。

實際的 CUDA 內(nèi)核非常簡單(如果您曾經(jīng)編程過 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巨型循環(huán)來執(zhí)行此操作，那么您會明白為什么這樣做會快得多。

使用訪問器

您可以在 CUDA 內(nèi)核中看到，我們直接處理正確類型的指針。 確實，直接在 cuda 內(nèi)核中使用高級類型不可知張量會非常低效。

但是，這是以易于使用和可讀性為代價的，尤其是對于高維數(shù)據(jù)。 在我們的示例中，例如，我們知道連續(xù)的gates張量具有 3 個維度：

1. 批次，batch_size的大小和3*state_size的步幅
2. 3的行，大小和state_size的步幅
3. 指數(shù)，state_size的大小和1的步幅

那么我們?nèi)绾卧L問內(nèi)核中的元素gates[n][row][column]？ 事實證明，您需要通過一些簡單的算法就可以大步訪問元素。

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

除了冗長之外，該表達式還需要跨步才能被明確地知道，并因此在其參數(shù)中傳遞給內(nèi)核函數(shù)。 您會看到，在內(nèi)核函數(shù)接受具有不同大小的多個張量的情況下，您將得到很長的參數(shù)列表。

對我們來說幸運的是，ATen 提供了通過動態(tài)檢查 Tensor 是尺寸的類型和數(shù)量而創(chuàng)建的訪問器。 然后，訪問器公開一個 API，可以有效地訪問 Tensor 元素，而不必轉(zhuǎn)換為單個指針：

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

訪問器對象具有較高級別的接口，具有.size()和.stride()方法以及多維索引。 .accessor<>接口旨在在 CPU 張量上有效訪問數(shù)據(jù)。 cuda 張量的等效項是packed_accessor64<>和packed_accessor32<>，它們產(chǎn)生具有 64 位或 32 位整數(shù)索引的壓縮訪問器。

與 Accessor 的根本區(qū)別在于，打包的 Accessor 在其結(jié)構(gòu)內(nèi)部復(fù)制大小和跨度數(shù)據(jù)，而不是指向它。 它允許我們將其傳遞給 CUDA 內(nèi)核函數(shù)并在其中使用其接口。

我們可以設(shè)計一個使用壓縮訪問器而不是指針的函數(shù)。

__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)

讓我們分解一下這里使用的模板。 前兩個參數(shù)scalar_t和2與常規(guī)訪問器相同。 參數(shù)torch::RestrictPtrTraits指示必須使用__restrict__關(guān)鍵字。 另請注意，我們使用了PackedAccessor32變體，將變體和步幅存儲在int32_t中。 這很重要，因為使用 64 位變體(PackedAccessor64）會使內(nèi)核變慢。

函數(shù)聲明變?yōu)?/p>

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

該實現(xiàn)更具可讀性！ 然后，通過在主機函數(shù)內(nèi)使用.packed_accessor32<>方法創(chuàng)建壓縮訪問器來調(diào)用此函數(shù)。

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_forward_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 的 op 與 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
    })

現(xiàn)在，我們使用CUDAExtension()代替CppExtension()。 我們只需要指定.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

現(xiàn)在使用我們的自定義 CUDA 內(nèi)核：

Forward: 129.431 us | Backward 304.641 us

更多性能提升！

結(jié)論

現(xiàn)在，您應(yīng)該已經(jīng)對 PyTorch 的 C ++擴展機制有了很好的了解，并有使用它們的動機。 您可以在此處找到本說明中顯示的代碼示例。 如有疑問，請使用論壇。 如果您遇到任何問題，也請務(wù)必查看我們的常見問題解答。