当前位置： 首页 > article >正文

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

article 2024/11/29 12:43:54

自定义 C++和 CUDA 扩展

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

译者：飞龙

协议：CC BY-NC-SA 4.0

作者：Peter Goldsborough

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

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

  • ATen 库是我们进行张量计算的主要 API，

  • pybind11 是我们为 C++ 代码创建 Python 绑定的方式，

  • 管理 ATen 和 pybind11 之间交互细节的头文件。

d_sigmoid()的实现展示了如何使用 ATen API。PyTorch 的张量和变量接口是从 ATen 库自动生成的，因此我们可以将我们的 Python 实现几乎一对一地转换成 C++。我们所有计算的主要数据类型将是torch::Tensor。其完整 API 可以在这里查看。还要注意，我们可以包含<iostream>或任何其他 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相同的信息。在后台，这将执行以下操作：

  1. 创建一个临时目录/tmp/torch_extensions/lltm，

  2. 在临时目录中生成一个Ninja构建文件，

  3. 将您的源文件编译成共享库，

  4. 将此共享库导入为 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 个维度：

  1. 批处理，batch_size的大小和3*state_size的步幅

  2. 行，3的大小和state_size的步幅

  3. 索引，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<float>()方法来获取底层数据的原始指针（就像之前 PyTorch 张量的.data_ptr<float>()一样）。我们还指定了张量的输出形状，我们将其硬编码为8 x 8。torch::from_blob的输出是一个指向 OpenCV 矩阵所拥有内存的torch::Tensor。

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

使用 TorchScript 注册自定义运算符

现在我们已经在 C++中实现了自定义运算符，我们需要在 TorchScript 运行时和编译器中注册它。这将允许 TorchScript 编译器解析 TorchScript 代码中对我们自定义运算符的引用。如果您曾经使用过 pybind11 库，我们的注册语法与 pybind11 语法非常相似。要注册单个函数，我们写入：

TORCH_LIBRARY(my_ops,  m)  {
  m.def("warp_perspective",  warp_perspective);
}

在我们的op.cpp文件的顶层某处。TORCH_LIBRARY宏创建一个函数，该函数在程序启动时将被调用。你的库的名称（my_ops）作为第一个参数给出（不应该用引号括起来）。第二个参数（m）定义了一个torch::Library类型的变量，它是注册你的运算符的主要接口。方法Library::def实际上创建了一个名为warp_perspective的运算符，将其暴露给 Python 和 TorchScript。你可以通过多次调用def来定义任意数量的运算符。

在幕后，def函数实际上做了很多工作：它使用模板元编程来检查函数的类型签名，并将其转换为一个运算符模式，该模式指定了 TorchScript 类型系统中的运算符类型。

构建自定义运算符

现在我们已经在 C++中实现了我们的自定义运算符并编写了其注册代码，是时候将运算符构建成一个（共享）库，以便我们可以将其加载到 Python 中进行研究和实验，或者加载到 C++中进行无 Python 环境中的推断。有多种方法可以构建我们的运算符，可以使用纯 CMake，也可以使用 Python 的替代方法，如setuptools。为简洁起见，以下段落仅讨论 CMake 方法。本教程的附录将深入探讨其他替代方法。

环境设置

我们需要安装 PyTorch 和 OpenCV。获取两者最简单和最独立于平台的方法是通过 Conda：

conda install -c pytorch pytorch
conda install opencv

使用 CMake 构建

使用CMake构建系统将我们的自定义运算符构建成一个共享库，我们需要编写一个简短的CMakeLists.txt文件，并将其与之前的op.cpp文件放在一起。为此，让我们同意一个看起来像这样的目录结构：

warp-perspective/
  op.cpp
  CMakeLists.txt

然后我们的CMakeLists.txt文件的内容应该是以下内容：

cmake_minimum_required(VERSION  3.1  FATAL_ERROR)
project(warp_perspective)

find_package(Torch  REQUIRED)
find_package(OpenCV  REQUIRED)

# Define our library target
add_library(warp_perspective  SHARED  op.cpp)
# Enable C++14
target_compile_features(warp_perspective  PRIVATE  cxx_std_14)
# Link against LibTorch
target_link_libraries(warp_perspective  "${TORCH_LIBRARIES}")
# Link against OpenCV
target_link_libraries(warp_perspective  opencv_core  opencv_imgproc)

现在要构建我们的运算符，我们可以从我们的warp_perspective文件夹中运行以下命令：

$  mkdir  build
$  cd  build
$  cmake  -DCMAKE_PREFIX_PATH="$(python  -c  'import torch.utils; print(torch.utils.cmake_prefix_path)')"  ..
--  The  C  compiler  identification  is  GNU  5.4.0
--  The  CXX  compiler  identification  is  GNU  5.4.0
--  Check  for  working  C  compiler:  /usr/bin/cc
--  Check  for  working  C  compiler:  /usr/bin/cc  --  works
--  Detecting  C  compiler  ABI  info
--  Detecting  C  compiler  ABI  info  -  done
--  Detecting  C  compile  features
--  Detecting  C  compile  features  -  done
--  Check  for  working  CXX  compiler:  /usr/bin/c++
--  Check  for  working  CXX  compiler:  /usr/bin/c++  --  works
--  Detecting  CXX  compiler  ABI  info
--  Detecting  CXX  compiler  ABI  info  -  done
--  Detecting  CXX  compile  features
--  Detecting  CXX  compile  features  -  done
--  Looking  for  pthread.h
--  Looking  for  pthread.h  -  found
--  Looking  for  pthread_create
--  Looking  for  pthread_create  -  not  found
--  Looking  for  pthread_create  in  pthreads
--  Looking  for  pthread_create  in  pthreads  -  not  found
--  Looking  for  pthread_create  in  pthread
--  Looking  for  pthread_create  in  pthread  -  found
--  Found  Threads:  TRUE
--  Found  torch:  /libtorch/lib/libtorch.so
--  Configuring  done
--  Generating  done
--  Build  files  have  been  written  to:  /warp_perspective/build
$  make  -j
Scanning  dependencies  of  target  warp_perspective
[  50%]  Building  CXX  object  CMakeFiles/warp_perspective.dir/op.cpp.o
[100%]  Linking  CXX  shared  library  libwarp_perspective.so
[100%]  Built  target  warp_perspective

这将在build文件夹中放置一个libwarp_perspective.so共享库文件。在上面的cmake命令中，我们使用辅助变量torch.utils.cmake_prefix_path方便地告诉我们 PyTorch 安装的 cmake 文件在哪里。

我们将在下面详细探讨如何使用和调用我们的运算符，但为了早点感受到成功，我们可以尝试在 Python 中运行以下代码：

import torch
torch.ops.load_library("build/libwarp_perspective.so")
print(torch.ops.my_ops.warp_perspective)

如果一切顺利，这应该打印出类似的内容：

<built-in method my_ops::warp_perspective of PyCapsule object at 0x7f618fc6fa50>

这是我们以后将用来调用我们自定义运算符的 Python 函数。

在 Python 中使用 TorchScript 自定义运算符

一旦我们的自定义运算符构建到一个共享库中，我们就可以在 Python 中的 TorchScript 模型中使用这个运算符。这有两个部分：首先将运算符加载到 Python 中，然后在 TorchScript 代码中使用该运算符。

你已经看到如何将你的运算符导入 Python：torch.ops.load_library()。这个函数接受包含自定义运算符的共享库路径，并将其加载到当前进程中。加载共享库还将执行TORCH_LIBRARY块。这将注册我们的自定义运算符到 TorchScript 编译器，并允许我们在 TorchScript 代码中使用该运算符。

你可以将加载的运算符称为torch.ops.<namespace>.<function>，其中<namespace>是你的运算符名称的命名空间部分，<function>是你的运算符的函数名称。对于我们上面编写的运算符，命名空间是my_ops，函数名称是warp_perspective，这意味着我们的运算符可以作为torch.ops.my_ops.warp_perspective使用。虽然这个函数可以在脚本化或跟踪的 TorchScript 模块中使用，我们也可以在普通的急切 PyTorch 中使用它，并传递常规的 PyTorch 张量：

import torch
torch.ops.load_library("build/libwarp_perspective.so")
print(torch.ops.my_ops.warp_perspective(torch.randn(32, 32), torch.rand(3, 3)))

生产：

tensor([[0.0000, 0.3218, 0.4611,  ..., 0.4636, 0.4636, 0.4636],
      [0.3746, 0.0978, 0.5005,  ..., 0.4636, 0.4636, 0.4636],
      [0.3245, 0.0169, 0.0000,  ..., 0.4458, 0.4458, 0.4458],
      ...,
      [0.1862, 0.1862, 0.1692,  ..., 0.0000, 0.0000, 0.0000],
      [0.1862, 0.1862, 0.1692,  ..., 0.0000, 0.0000, 0.0000],
      [0.1862, 0.1862, 0.1692,  ..., 0.0000, 0.0000, 0.0000]])

注意

在幕后发生的事情是，当您在 Python 中第一次访问torch.ops.namespace.function时，TorchScript 编译器（在 C++领域）将查看是否已经注册了函数namespace::function，如果是，则返回一个 Python 句柄到这个函数，我们随后可以使用这个句柄从 Python 调用我们的 C++运算符实现。这是 TorchScript 自定义运算符和 C++扩展之间的一个值得注意的区别：C++扩展是通过 pybind11 手动绑定的，而 TorchScript 自定义运算符是由 PyTorch 自身动态绑定的。Pybind11 在绑定到 Python 时给您更多的灵活性，因此建议用于纯粹的急切代码，但不支持 TorchScript 运算符。

从这里开始，您可以在脚本化或跟踪的代码中像使用torch包中的其他函数一样使用您的自定义运算符。事实上，“标准库”函数如torch.matmul通过与自定义运算符基本相同的注册路径，这使得自定义运算符在 TorchScript 中如何以及在哪里使用时成为真正的一等公民。（然而，一个区别是，标准库函数具有自定义编写的 Python 参数解析逻辑，与torch.ops参数解析不同。）

使用跟踪的自定义运算符

让我们首先将我们的运算符嵌入到一个跟踪函数中。回想一下，对于跟踪，我们从一些普通的 PyTorch 代码开始：

def compute(x, y, z):
    return x.matmul(y) + torch.relu(z)

然后在其上调用torch.jit.trace。我们进一步传递给torch.jit.trace一些示例输入，它将转发给我们的实现以记录输入流经过时发生的操作序列。这样做的结果实际上是急切 PyTorch 程序的“冻结”版本，TorchScript 编译器可以进一步分析、优化和序列化：

inputs = [torch.randn(4, 8), torch.randn(8, 5), torch.randn(4, 5)]
trace = torch.jit.trace(compute, inputs)
print(trace.graph)

生成：

graph(%x : Float(4:8, 8:1),
      %y : Float(8:5, 5:1),
      %z : Float(4:5, 5:1)):
  %3 : Float(4:5, 5:1) = aten::matmul(%x, %y) # test.py:10:0
  %4 : Float(4:5, 5:1) = aten::relu(%z) # test.py:10:0
  %5 : int = prim::Constant[value=1]() # test.py:10:0
  %6 : Float(4:5, 5:1) = aten::add(%3, %4, %5) # test.py:10:0
  return (%6)

现在，令人兴奋的发现是，我们可以简单地将我们的自定义运算符放入我们的 PyTorch 跟踪中，就像它是torch.relu或任何其他torch函数一样：

def compute(x, y, z):
    x = torch.ops.my_ops.warp_perspective(x, torch.eye(3))
    return x.matmul(y) + torch.relu(z)

然后像以前一样对其进行跟踪：

inputs = [torch.randn(4, 8), torch.randn(8, 5), torch.randn(8, 5)]
trace = torch.jit.trace(compute, inputs)
print(trace.graph)

生成：

graph(%x.1 : Float(4:8, 8:1),
      %y : Float(8:5, 5:1),
      %z : Float(8:5, 5:1)):
  %3 : int = prim::Constant[value=3]() # test.py:25:0
  %4 : int = prim::Constant[value=6]() # test.py:25:0
  %5 : int = prim::Constant[value=0]() # test.py:25:0
  %6 : Device = prim::Constant[value="cpu"]() # test.py:25:0
  %7 : bool = prim::Constant[value=0]() # test.py:25:0
  %8 : Float(3:3, 3:1) = aten::eye(%3, %4, %5, %6, %7) # test.py:25:0
  %x : Float(8:8, 8:1) = my_ops::warp_perspective(%x.1, %8) # test.py:25:0
  %10 : Float(8:5, 5:1) = aten::matmul(%x, %y) # test.py:26:0
  %11 : Float(8:5, 5:1) = aten::relu(%z) # test.py:26:0
  %12 : int = prim::Constant[value=1]() # test.py:26:0
  %13 : Float(8:5, 5:1) = aten::add(%10, %11, %12) # test.py:26:0
  return (%13)

将 TorchScript 自定义运算符集成到跟踪的 PyTorch 代码中就像这样简单！

使用脚本的自定义运算符

除了跟踪之外，另一种获得 PyTorch 程序的 TorchScript 表示的方法是直接在 TorchScript 中编写代码。TorchScript 在很大程度上是 Python 语言的一个子集，具有一些限制，使得 TorchScript 编译器更容易推理程序。通过使用@torch.jit.script对自由函数进行注释，以及对类中的方法使用@torch.jit.script_method（该类还必须派生自torch.jit.ScriptModule），您可以将常规的 PyTorch 代码转换为 TorchScript。有关 TorchScript 注释的更多详细信息，请参见这里。

使用 TorchScript 而不是跟踪的一个特别原因是，跟踪无法捕获 PyTorch 代码中的控制流。因此，让我们考虑这个使用控制流的函数：

def compute(x, y):
  if bool(x[0][0] == 42):
      z = 5
  else:
      z = 10
  return x.matmul(y) + z

要将这个函数从普通的 PyTorch 转换为 TorchScript，我们使用@torch.jit.script对其进行注释：

@torch.jit.script
def compute(x, y):
  if bool(x[0][0] == 42):
      z = 5
  else:
      z = 10
  return x.matmul(y) + z

这将把compute函数即时编译成图表示，我们可以在compute.graph属性中检查它：

>>> compute.graph
graph(%x : Dynamic
 %y : Dynamic) {
 %14 : int = prim::Constant[value=1]()
 %2 : int = prim::Constant[value=0]()
 %7 : int = prim::Constant[value=42]()
 %z.1 : int = prim::Constant[value=5]()
 %z.2 : int = prim::Constant[value=10]()
 %4 : Dynamic = aten::select(%x, %2, %2)
 %6 : Dynamic = aten::select(%4, %2, %2)
 %8 : Dynamic = aten::eq(%6, %7)
 %9 : bool = prim::TensorToBool(%8)
 %z : int = prim::If(%9)
 block0() {
 -> (%z.1)
 }
 block1() {
 -> (%z.2)
 }
 %13 : Dynamic = aten::matmul(%x, %y)
 %15 : Dynamic = aten::add(%13, %z, %14)
 return (%15);
}

现在，就像以前一样，我们可以在我们的脚本代码中像使用任何其他函数一样使用我们的自定义运算符：

torch.ops.load_library("libwarp_perspective.so")

@torch.jit.script
def compute(x, y):
  if bool(x[0] == 42):
      z = 5
  else:
      z = 10
  x = torch.ops.my_ops.warp_perspective(x, torch.eye(3))
  return x.matmul(y) + z

当 TorchScript 编译器看到对torch.ops.my_ops.warp_perspective的引用时，它将找到我们通过 C++中的TORCH_LIBRARY函数注册的实现，并将其编译成其图表示：

>>> compute.graph
graph(%x.1 : Dynamic
 %y : Dynamic) {
 %20 : int = prim::Constant[value=1]()
 %16 : int[] = prim::Constant[value=[0, -1]]()
 %14 : int = prim::Constant[value=6]()
 %2 : int = prim::Constant[value=0]()
 %7 : int = prim::Constant[value=42]()
 %z.1 : int = prim::Constant[value=5]()
 %z.2 : int = prim::Constant[value=10]()
 %13 : int = prim::Constant[value=3]()
 %4 : Dynamic = aten::select(%x.1, %2, %2)
 %6 : Dynamic = aten::select(%4, %2, %2)
 %8 : Dynamic = aten::eq(%6, %7)
 %9 : bool = prim::TensorToBool(%8)
 %z : int = prim::If(%9)
 block0() {
 -> (%z.1)
 }
 block1() {
 -> (%z.2)
 }
 %17 : Dynamic = aten::eye(%13, %14, %2, %16)
 %x : Dynamic = my_ops::warp_perspective(%x.1, %17)
 %19 : Dynamic = aten::matmul(%x, %y)
 %21 : Dynamic = aten::add(%19, %z, %20)
 return (%21);
 }

特别注意图的末尾对my_ops::warp_perspective的引用。

注意

TorchScript 图表示仍然可能会发生变化。不要依赖它看起来像这样。

这就是在 Python 中使用我们的自定义运算符时的全部内容。简而言之，您可以使用torch.ops.load_library导入包含您的运算符的库，并像从您的跟踪或脚本化的 TorchScript 代码中调用任何其他torch运算符一样调用您的自定义运算符。

在 C++中使用 TorchScript 自定义运算符

TorchScript 的一个有用功能是将模型序列化为磁盘文件。这个文件可以通过网络发送，存储在文件系统中，或者更重要的是，可以在不需要保留原始源代码的情况下动态反序列化和执行。这在 Python 中是可能的，但在 C++中也是可能的。为此，PyTorch 提供了一个纯 C++ API用于反序列化以及执行 TorchScript 模型。如果您还没有，请阅读在 C++中加载和运行序列化的 TorchScript 模型的教程，接下来的几段将基于此构建。

简而言之，即使从文件中反序列化并在 C++中运行，自定义运算符也可以像常规的torch运算符一样执行。这唯一的要求是将我们之前构建的自定义运算符共享库与我们在其中执行模型的 C++应用程序链接起来。在 Python 中，这只需简单调用torch.ops.load_library。在 C++中，您需要将共享库与您正在使用的任何构建系统中的主应用程序链接起来。以下示例将使用 CMake 展示这一点。

注意

从技术上讲，您也可以在运行时以与我们在 Python 中所做的方式相同的方式动态加载共享库到您的 C++应用程序中。在 Linux 上，您可以使用 dlopen 来做到这一点。其他平台上也存在等价物。

在上面链接的 C++执行教程的基础上，让我们从一个最小的 C++应用程序开始，该应用程序位于一个不同的文件夹中的main.cpp文件中，加载并执行一个序列化的 TorchScript 模型：

#include  <torch/script.h> // One-stop header.

#include  <iostream>
#include  <memory>

int  main(int  argc,  const  char*  argv[])  {
  if  (argc  !=  2)  {
  std::cerr  <<  "usage: example-app <path-to-exported-script-module>\n";
  return  -1;
  }

  // Deserialize the ScriptModule from a file using torch::jit::load().
  torch::jit::script::Module  module  =  torch::jit::load(argv[1]);

  std::vector<torch::jit::IValue>  inputs;
  inputs.push_back(torch::randn({4,  8}));
  inputs.push_back(torch::randn({8,  5}));

  torch::Tensor  output  =  module.forward(std::move(inputs)).toTensor();

  std::cout  <<  output  <<  std::endl;
}

还有一个小的CMakeLists.txt文件：

cmake_minimum_required(VERSION  3.1  FATAL_ERROR)
project(example_app)

find_package(Torch  REQUIRED)

add_executable(example_app  main.cpp)
target_link_libraries(example_app  "${TORCH_LIBRARIES}")
target_compile_features(example_app  PRIVATE  cxx_range_for)

在这一点上，我们应该能够构建应用程序：

$  mkdir  build
$  cd  build
$  cmake  -DCMAKE_PREFIX_PATH="$(python  -c  'import torch.utils; print(torch.utils.cmake_prefix_path)')"  ..
--  The  C  compiler  identification  is  GNU  5.4.0
--  The  CXX  compiler  identification  is  GNU  5.4.0
--  Check  for  working  C  compiler:  /usr/bin/cc
--  Check  for  working  C  compiler:  /usr/bin/cc  --  works
--  Detecting  C  compiler  ABI  info
--  Detecting  C  compiler  ABI  info  -  done
--  Detecting  C  compile  features
--  Detecting  C  compile  features  -  done
--  Check  for  working  CXX  compiler:  /usr/bin/c++
--  Check  for  working  CXX  compiler:  /usr/bin/c++  --  works
--  Detecting  CXX  compiler  ABI  info
--  Detecting  CXX  compiler  ABI  info  -  done
--  Detecting  CXX  compile  features
--  Detecting  CXX  compile  features  -  done
--  Looking  for  pthread.h
--  Looking  for  pthread.h  -  found
--  Looking  for  pthread_create
--  Looking  for  pthread_create  -  not  found
--  Looking  for  pthread_create  in  pthreads
--  Looking  for  pthread_create  in  pthreads  -  not  found
--  Looking  for  pthread_create  in  pthread
--  Looking  for  pthread_create  in  pthread  -  found
--  Found  Threads:  TRUE
--  Found  torch:  /libtorch/lib/libtorch.so
--  Configuring  done
--  Generating  done
--  Build  files  have  been  written  to:  /example_app/build
$  make  -j
Scanning  dependencies  of  target  example_app
[  50%]  Building  CXX  object  CMakeFiles/example_app.dir/main.cpp.o
[100%]  Linking  CXX  executable  example_app
[100%]  Built  target  example_app

并且在不传递模型的情况下运行它：

$  ./example_app
usage:  example_app  <path-to-exported-script-module>

接下来，让我们序列化我们之前编写的使用我们自定义运算符的脚本函数：

torch.ops.load_library("libwarp_perspective.so")

@torch.jit.script
def compute(x, y):
  if bool(x[0][0] == 42):
      z = 5
  else:
      z = 10
  x = torch.ops.my_ops.warp_perspective(x, torch.eye(3))
  return x.matmul(y) + z

compute.save("example.pt")

最后一行将脚本函数序列化为一个名为“example.pt”的文件。如果我们将这个序列化模型传递给我们的 C++应用程序，我们可以立即运行它：

$  ./example_app  example.pt
terminate  called  after  throwing  an  instance  of  'torch::jit::script::ErrorReport'
what():
Schema  not  found  for  node.  File  a  bug  report.
Node:  %16  :  Dynamic  =  my_ops::warp_perspective(%0,  %19)

或者也许不是。也许还不是。当然！我们还没有将自定义运算符库与我们的应用程序链接起来。让我们立即做这个，为了正确地做这件事，让我们稍微更新我们的文件组织，看起来像这样：

example_app/
  CMakeLists.txt
  main.cpp
  warp_perspective/
    CMakeLists.txt
    op.cpp

这将允许我们将warp_perspective库 CMake 目标作为我们应用程序目标的子目录。example_app文件夹中的顶层CMakeLists.txt应该如下所示：

cmake_minimum_required(VERSION  3.1  FATAL_ERROR)
project(example_app)

find_package(Torch  REQUIRED)

add_subdirectory(warp_perspective)

add_executable(example_app  main.cpp)
target_link_libraries(example_app  "${TORCH_LIBRARIES}")
target_link_libraries(example_app  -Wl,--no-as-needed  warp_perspective)
target_compile_features(example_app  PRIVATE  cxx_range_for)

这个基本的 CMake 配置看起来与以前很像，只是我们将warp_perspective CMake 构建添加为一个子目录。一旦它的 CMake 代码运行，我们将我们的example_app应用程序与warp_perspective共享库链接起来。

注意

上面示例中嵌入了一个关键细节：warp_perspective链接行前缀-Wl,--no-as-needed。这是必需的，因为我们实际上不会在应用程序代码中调用warp_perspective共享库中的任何函数。我们只需要TORCH_LIBRARY函数运行。不方便的是，这会让链接器混淆，并使其认为可以完全跳过与库的链接。在 Linux 上，-Wl,--no-as-needed标志强制进行链接（注意：此标志特定于 Linux！）。还有其他解决方法。最简单的方法是在您需要从主应用程序调用的运算符库中定义某个函数。这可以是在某个头文件中声明的简单函数void init();，然后在运算符库中定义为void init() { }。在主应用程序中调用此init()函数将使链接器认为这是值得链接的库。不幸的是，这超出了我们的控制范围，我们宁愿让您了解这个原因和简单的解决方法，而不是给您一些不透明的宏来放入您的代码中。

现在，由于我们现在在顶层找到了Torch包，warp_perspective子目录中的CMakeLists.txt文件可以稍微缩短一点。它应该是这样的：

find_package(OpenCV  REQUIRED)
add_library(warp_perspective  SHARED  op.cpp)
target_compile_features(warp_perspective  PRIVATE  cxx_range_for)
target_link_libraries(warp_perspective  PRIVATE  "${TORCH_LIBRARIES}")
target_link_libraries(warp_perspective  PRIVATE  opencv_core  opencv_photo)

让我们重新构建我们的示例应用程序，它还将链接到自定义运算符库。在顶层example_app目录中：

$  mkdir  build
$  cd  build
$  cmake  -DCMAKE_PREFIX_PATH="$(python  -c  'import torch.utils; print(torch.utils.cmake_prefix_path)')"  ..
--  The  C  compiler  identification  is  GNU  5.4.0
--  The  CXX  compiler  identification  is  GNU  5.4.0
--  Check  for  working  C  compiler:  /usr/bin/cc
--  Check  for  working  C  compiler:  /usr/bin/cc  --  works
--  Detecting  C  compiler  ABI  info
--  Detecting  C  compiler  ABI  info  -  done
--  Detecting  C  compile  features
--  Detecting  C  compile  features  -  done
--  Check  for  working  CXX  compiler:  /usr/bin/c++
--  Check  for  working  CXX  compiler:  /usr/bin/c++  --  works
--  Detecting  CXX  compiler  ABI  info
--  Detecting  CXX  compiler  ABI  info  -  done
--  Detecting  CXX  compile  features
--  Detecting  CXX  compile  features  -  done
--  Looking  for  pthread.h
--  Looking  for  pthread.h  -  found
--  Looking  for  pthread_create
--  Looking  for  pthread_create  -  not  found
--  Looking  for  pthread_create  in  pthreads
--  Looking  for  pthread_create  in  pthreads  -  not  found
--  Looking  for  pthread_create  in  pthread
--  Looking  for  pthread_create  in  pthread  -  found
--  Found  Threads:  TRUE
--  Found  torch:  /libtorch/lib/libtorch.so
--  Configuring  done
--  Generating  done
--  Build  files  have  been  written  to:  /warp_perspective/example_app/build
$  make  -j
Scanning  dependencies  of  target  warp_perspective
[  25%]  Building  CXX  object  warp_perspective/CMakeFiles/warp_perspective.dir/op.cpp.o
[  50%]  Linking  CXX  shared  library  libwarp_perspective.so
[  50%]  Built  target  warp_perspective
Scanning  dependencies  of  target  example_app
[  75%]  Building  CXX  object  CMakeFiles/example_app.dir/main.cpp.o
[100%]  Linking  CXX  executable  example_app
[100%]  Built  target  example_app

如果我们现在运行example_app二进制文件并将序列化模型交给它，我们应该会得到一个美好的结局：

$  ./example_app  example.pt
11.4125  5.8262  9.5345  8.6111  12.3997
  7.4683  13.5969  9.0850  11.0698  9.4008
  7.4597  15.0926  12.5727  8.9319  9.0666
  9.4834  11.1747  9.0162  10.9521  8.6269
10.0000  10.0000  10.0000  10.0000  10.0000
10.0000  10.0000  10.0000  10.0000  10.0000
10.0000  10.0000  10.0000  10.0000  10.0000
10.0000  10.0000  10.0000  10.0000  10.0000
[  Variable[CPUFloatType]{8,5}  ]

成功！您现在已经准备好进行推理了。

结论

本教程向您展示了如何在 C++中实现自定义 TorchScript 运算符，如何将其构建为共享库，如何在 Python 中使用它来定义 TorchScript 模型，最后如何将其加载到用于推理工作负载的 C++应用程序中。您现在已经准备好通过 C++运算符扩展您的 TorchScript 模型，这些运算符与第三方 C++库进行接口，编写自定义高性能 CUDA 内核，或实现任何其他需要 Python、TorchScript 和 C++之间无缝融合的用例。

如往常一样，如果遇到任何问题或有疑问，您可以使用我们的论坛或GitHub 问题联系我们。此外，我们的常见问题（FAQ）页面可能会提供有用的信息。

附录 A：构建自定义运算符的更多方法

“构建自定义运算符”部分解释了如何使用 CMake 将自定义运算符构建为共享库。本附录概述了两种进一步的编译方法。它们都使用 Python 作为编译过程的“驱动程序”或“接口”。此外，它们都重用了 PyTorch 为C++扩展提供的现有基础设施，这些扩展是依赖于pybind11的 TorchScript 自定义运算符的等效版本，用于将 C++函数“显式”绑定到 Python 中。

第一种方法使用 C++扩展的方便的即时（JIT）编译接口在您首次运行 PyTorch 脚本时在后台编译您的代码。第二种方法依赖于古老的setuptools包，并涉及编写一个单独的setup.py文件。这允许更高级的配置以及与其他基于setuptools的项目集成。我们将在下面详细探讨这两种方法。

使用 JIT 编译进行构建

PyTorch C++扩展工具包提供的 JIT 编译功能允许将自定义运算符的编译直接嵌入到您的 Python 代码中，例如在您的训练脚本顶部。

注意

这里的“JIT 编译”与 TorchScript 编译器中进行的 JIT 编译优化程序无关。它只是意味着您的自定义运算符 C++代码将在您首次导入时编译到系统的/tmp 目录下的一个文件夹中，就好像您之前自己编译过一样。

这个 JIT 编译功能有两种方式。在第一种方式中，您仍然将您的运算符实现放在一个单独的文件中（op.cpp），然后使用torch.utils.cpp_extension.load()来编译您的扩展。通常，这个函数会返回暴露您的 C++扩展的 Python 模块。然而，由于我们没有将自定义运算符编译成自己的 Python 模块，我们只想编译一个普通的共享库。幸运的是，torch.utils.cpp_extension.load()有一个参数is_python_module，我们可以将其设置为False，以指示我们只对构建共享库感兴趣，而不是 Python 模块。torch.utils.cpp_extension.load()然后会编译并加载共享库到当前进程中，就像之前torch.ops.load_library做的那样：

import torch.utils.cpp_extension

torch.utils.cpp_extension.load(
    name="warp_perspective",
    sources=["op.cpp"],
    extra_ldflags=["-lopencv_core", "-lopencv_imgproc"],
    is_python_module=False,
    verbose=True
)

print(torch.ops.my_ops.warp_perspective)

这应该大致打印：

<built-in method my_ops::warp_perspective of PyCapsule object at 0x7f3e0f840b10>

第二种 JIT 编译的方式允许您将自定义 TorchScript 运算符的源代码作为字符串传递。为此，请使用torch.utils.cpp_extension.load_inline：

import torch
import torch.utils.cpp_extension

op_source = """
#include <opencv2/opencv.hpp>
#include <torch/script.h>

torch::Tensor warp_perspective(torch::Tensor image, torch::Tensor warp) {
 cv::Mat image_mat(/*rows=*/image.size(0),
 /*cols=*/image.size(1),
 /*type=*/CV_32FC1,
 /*data=*/image.data<float>());
 cv::Mat warp_mat(/*rows=*/warp.size(0),
 /*cols=*/warp.size(1),
 /*type=*/CV_32FC1,
 /*data=*/warp.data<float>());

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

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

TORCH_LIBRARY(my_ops, m) {
 m.def("warp_perspective", &warp_perspective);
}
"""

torch.utils.cpp_extension.load_inline(
    name="warp_perspective",
    cpp_sources=op_source,
    extra_ldflags=["-lopencv_core", "-lopencv_imgproc"],
    is_python_module=False,
    verbose=True,
)

print(torch.ops.my_ops.warp_perspective)

自然地，最佳实践是只在您的源代码相当短的情况下使用torch.utils.cpp_extension.load_inline。

请注意，如果您在 Jupyter Notebook 中使用这个功能，不要多次执行注册单元格，因为每次执行都会注册一个新的库并重新注册自定义运算符。如果需要重新执行，请在此之前重新启动笔记本的 Python 内核。

使用 Setuptools 构建

从 Python 中构建我们的自定义运算符的第二种方法是使用setuptools。这样做的好处是setuptools具有一个非常强大和广泛的接口，用于构建用 C++编写的 Python 模块。然而，由于setuptools实际上是用于构建 Python 模块而不是普通的共享库（这些库没有模块所需的入口点），这条路线可能有点古怪。也就是说，您只需要一个setup.py文件来替代CMakeLists.txt，它看起来像这样：

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension

setup(
    name="warp_perspective",
    ext_modules=[
        CppExtension(
            "warp_perspective",
            ["example_app/warp_perspective/op.cpp"],
            libraries=["opencv_core", "opencv_imgproc"],
        )
    ],
    cmdclass={"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)},
)

请注意，在底部的BuildExtension中启用了no_python_abi_suffix选项。这指示setuptools在生成的共享库名称中省略任何 Python-3 特定的 ABI 后缀。否则，在 Python 3.7 中，库可能被称为warp_perspective.cpython-37m-x86_64-linux-gnu.so，其中cpython-37m-x86_64-linux-gnu是 ABI 标签，但我们真的只想让它被称为warp_perspective.so。

如果我们现在在包含setup.py的文件夹中的终端中运行python setup.py build develop，我们应该会看到类似以下的内容：

$  python  setup.py  build  develop
running  build
running  build_ext
building  'warp_perspective'  extension
creating  build
creating  build/temp.linux-x86_64-3.7
gcc  -pthread  -B  /root/local/miniconda/compiler_compat  -Wl,--sysroot=/  -Wsign-compare  -DNDEBUG  -g  -fwrapv  -O3  -Wall  -Wstrict-prototypes  -fPIC  -I/root/local/miniconda/lib/python3.7/site-packages/torch/lib/include  -I/root/local/miniconda/lib/python3.7/site-packages/torch/lib/include/torch/csrc/api/include  -I/root/local/miniconda/lib/python3.7/site-packages/torch/lib/include/TH  -I/root/local/miniconda/lib/python3.7/site-packages/torch/lib/include/THC  -I/root/local/miniconda/include/python3.7m  -c  op.cpp  -o  build/temp.linux-x86_64-3.7/op.o  -DTORCH_API_INCLUDE_EXTENSION_H  -DTORCH_EXTENSION_NAME=warp_perspective  -D_GLIBCXX_USE_CXX11_ABI=0  -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  /root/local/miniconda/compiler_compat  -L/root/local/miniconda/lib  -Wl,-rpath=/root/local/miniconda/lib  -Wl,--no-as-needed  -Wl,--sysroot=/  build/temp.linux-x86_64-3.7/op.o  -lopencv_core  -lopencv_imgproc  -o  build/lib.linux-x86_64-3.7/warp_perspective.so
running  develop
running  egg_info
creating  warp_perspective.egg-info
writing  warp_perspective.egg-info/PKG-INFO
writing  dependency_links  to  warp_perspective.egg-info/dependency_links.txt
writing  top-level  names  to  warp_perspective.egg-info/top_level.txt
writing  manifest  file  'warp_perspective.egg-info/SOURCES.txt'
reading  manifest  file  'warp_perspective.egg-info/SOURCES.txt'
writing  manifest  file  'warp_perspective.egg-info/SOURCES.txt'
running  build_ext
copying  build/lib.linux-x86_64-3.7/warp_perspective.so  ->
Creating  /root/local/miniconda/lib/python3.7/site-packages/warp-perspective.egg-link  (link  to  .)
Adding  warp-perspective  0.0.0  to  easy-install.pth  file

Installed  /warp_perspective
Processing  dependencies  for  warp-perspective==0.0.0
Finished  processing  dependencies  for  warp-perspective==0.0.0

这将生成一个名为warp_perspective.so的共享库，我们可以像之前那样将其传递给torch.ops.load_library，以使我们的运算符对 TorchScript 可见：

>>> import torch
>>> torch.ops.load_library("warp_perspective.so")
>>> print(torch.ops.my_ops.warp_perspective)
<built-in method custom::warp_perspective of PyCapsule object at 0x7ff51c5b7bd0>

使用自定义 C++类扩展 TorchScript

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

译者：飞龙

协议：CC BY-NC-SA 4.0

本教程是自定义运算符教程的后续，介绍了我们为将 C++类绑定到 TorchScript 和 Python 中构建的 API。该 API 与pybind11非常相似，如果您熟悉该系统，大部分概念都会转移到这里。

在 C++中实现和绑定类

在本教程中，我们将定义一个简单的 C++类，该类在成员变量中维护持久状态。

// This header is all you need to do the C++ portions of this
// tutorial
#include  <torch/script.h>
// This header is what defines the custom class registration
// behavior specifically. script.h already includes this, but
// we include it here so you know it exists in case you want
// to look at the API or implementation.
#include  <torch/custom_class.h>

#include  <string>
#include  <vector>

template  <class  T>
struct  MyStackClass  :  torch::CustomClassHolder  {
  std::vector<T>  stack_;
  MyStackClass(std::vector<T>  init)  :  stack_(init.begin(),  init.end())  {}

  void  push(T  x)  {
  stack_.push_back(x);
  }
  T  pop()  {
  auto  val  =  stack_.back();
  stack_.pop_back();
  return  val;
  }

  c10::intrusive_ptr<MyStackClass>  clone()  const  {
  return  c10::make_intrusive<MyStackClass>(stack_);
  }

  void  merge(const  c10::intrusive_ptr<MyStackClass>&  c)  {
  for  (auto&  elem  :  c->stack_)  {
  push(elem);
  }
  }
};

有几点需要注意：

  • torch/custom_class.h是您需要包含的头文件，以便使用自定义类扩展 TorchScript。

  • 请注意，每当我们使用自定义类的实例时，我们都是通过c10::intrusive_ptr<>的实例来进行的。将intrusive_ptr视为类似于std::shared_ptr的智能指针，但引用计数直接存储在对象中，而不是存储在单独的元数据块中（就像在std::shared_ptr中所做的那样）。torch::Tensor内部使用相同的指针类型；自定义类也必须使用这种指针类型，以便我们可以一致地管理不同的对象类型。

  • 第二件要注意的事情是，用户定义的类必须继承自torch::CustomClassHolder。这确保了自定义类有空间来存储引用计数。

现在让我们看看如何使这个类对 TorchScript 可见，这个过程称为绑定类：

// Notice a few things:
// - We pass the class to be registered as a template parameter to
//   `torch::class_`. In this instance, we've passed the
//   specialization of the MyStackClass class ``MyStackClass<std::string>``.
//   In general, you cannot register a non-specialized template
//   class. For non-templated classes, you can just pass the
//   class name directly as the template parameter.
// - The arguments passed to the constructor make up the "qualified name"
//   of the class. In this case, the registered class will appear in
//   Python and C++ as `torch.classes.my_classes.MyStackClass`. We call
//   the first argument the "namespace" and the second argument the
//   actual class name.
TORCH_LIBRARY(my_classes,  m)  {
  m.class_<MyStackClass<std::string>>("MyStackClass")
  // The following line registers the contructor of our MyStackClass
  // class that takes a single `std::vector<std::string>` argument,
  // i.e. it exposes the C++ method `MyStackClass(std::vector<T> init)`.
  // Currently, we do not support registering overloaded
  // constructors, so for now you can only `def()` one instance of
  // `torch::init`.
  .def(torch::init<std::vector<std::string>>())
  // The next line registers a stateless (i.e. no captures) C++ lambda
  // function as a method. Note that a lambda function must take a
  // `c10::intrusive_ptr<YourClass>` (or some const/ref version of that)
  // as the first argument. Other arguments can be whatever you want.
  .def("top",  [](const  c10::intrusive_ptr<MyStackClass<std::string>>&  self)  {
  return  self->stack_.back();
  })
  // The following four lines expose methods of the MyStackClass<std::string>
  // class as-is. `torch::class_` will automatically examine the
  // argument and return types of the passed-in method pointers and
  // expose these to Python and TorchScript accordingly. Finally, notice
  // that we must take the *address* of the fully-qualified method name,
  // i.e. use the unary `&` operator, due to C++ typing rules.
  .def("push",  &MyStackClass<std::string>::push)
  .def("pop",  &MyStackClass<std::string>::pop)
  .def("clone",  &MyStackClass<std::string>::clone)
  .def("merge",  &MyStackClass<std::string>::merge)
  ;
}

使用 CMake 将示例构建为 C++项目

现在，我们将使用CMake构建系统构建上述 C++代码。首先，将我们迄今为止涵盖的所有 C++代码放入一个名为class.cpp的文件中。然后，编写一个简单的CMakeLists.txt文件并将其放在同一目录中。CMakeLists.txt应该如下所示：

cmake_minimum_required(VERSION  3.1  FATAL_ERROR)
project(custom_class)

find_package(Torch  REQUIRED)

# Define our library target
add_library(custom_class  SHARED  class.cpp)
set(CMAKE_CXX_STANDARD  14)
# Link against LibTorch
target_link_libraries(custom_class  "${TORCH_LIBRARIES}")

同时，创建一个build目录。您的文件树应该如下所示：

custom_class_project/
  class.cpp
  CMakeLists.txt
  build/

我们假设您已经按照上一个教程中描述的方式设置了环境。继续调用 cmake，然后进行构建：

$  cd  build
$  cmake  -DCMAKE_PREFIX_PATH="$(python  -c  'import torch.utils; print(torch.utils.cmake_prefix_path)')"  ..
  --  The  C  compiler  identification  is  GNU  7.3.1
  --  The  CXX  compiler  identification  is  GNU  7.3.1
  --  Check  for  working  C  compiler:  /opt/rh/devtoolset-7/root/usr/bin/cc
  --  Check  for  working  C  compiler:  /opt/rh/devtoolset-7/root/usr/bin/cc  --  works
  --  Detecting  C  compiler  ABI  info
  --  Detecting  C  compiler  ABI  info  -  done
  --  Detecting  C  compile  features
  --  Detecting  C  compile  features  -  done
  --  Check  for  working  CXX  compiler:  /opt/rh/devtoolset-7/root/usr/bin/c++
  --  Check  for  working  CXX  compiler:  /opt/rh/devtoolset-7/root/usr/bin/c++  --  works
  --  Detecting  CXX  compiler  ABI  info
  --  Detecting  CXX  compiler  ABI  info  -  done
  --  Detecting  CXX  compile  features
  --  Detecting  CXX  compile  features  -  done
  --  Looking  for  pthread.h
  --  Looking  for  pthread.h  -  found
  --  Looking  for  pthread_create
  --  Looking  for  pthread_create  -  not  found
  --  Looking  for  pthread_create  in  pthreads
  --  Looking  for  pthread_create  in  pthreads  -  not  found
  --  Looking  for  pthread_create  in  pthread
  --  Looking  for  pthread_create  in  pthread  -  found
  --  Found  Threads:  TRUE
  --  Found  torch:  /torchbind_tutorial/libtorch/lib/libtorch.so
  --  Configuring  done
  --  Generating  done
  --  Build  files  have  been  written  to:  /torchbind_tutorial/build
$  make  -j
  Scanning  dependencies  of  target  custom_class
  [  50%]  Building  CXX  object  CMakeFiles/custom_class.dir/class.cpp.o
  [100%]  Linking  CXX  shared  library  libcustom_class.so
  [100%]  Built  target  custom_class

您会发现现在（除其他内容外）在构建目录中存在一个动态库文件。在 Linux 上，这个文件可能被命名为libcustom_class.so。因此，文件树应该如下所示：

custom_class_project/
  class.cpp
  CMakeLists.txt
  build/
    libcustom_class.so

从 Python 和 TorchScript 中使用 C++类

现在，我们已经将我们的类及其注册编译到一个.so文件中，我们可以将该.so 加载到 Python 中并尝试它。以下是演示这一点的脚本：

import torch

# `torch.classes.load_library()` allows you to pass the path to your .so file
# to load it in and make the custom C++ classes available to both Python and
# TorchScript
torch.classes.load_library("build/libcustom_class.so")
# You can query the loaded libraries like this:
print(torch.classes.loaded_libraries)
# prints {'/custom_class_project/build/libcustom_class.so'}

# We can find and instantiate our custom C++ class in python by using the
# `torch.classes` namespace:
#
# This instantiation will invoke the MyStackClass(std::vector<T> init)
# constructor we registered earlier
s = torch.classes.my_classes.MyStackClass(["foo", "bar"])

# We can call methods in Python
s.push("pushed")
assert s.pop() == "pushed"

# Test custom operator
s.push("pushed")
torch.ops.my_classes.manipulate_instance(s)  # acting as s.pop()
assert s.top() == "bar" 

# Returning and passing instances of custom classes works as you'd expect
s2 = s.clone()
s.merge(s2)
for expected in ["bar", "foo", "bar", "foo"]:
    assert s.pop() == expected

# We can also use the class in TorchScript
# For now, we need to assign the class's type to a local in order to
# annotate the type on the TorchScript function. This may change
# in the future.
MyStackClass = torch.classes.my_classes.MyStackClass

@torch.jit.script
def do_stacks(s: MyStackClass):  # We can pass a custom class instance
    # We can instantiate the class
    s2 = torch.classes.my_classes.MyStackClass(["hi", "mom"])
    s2.merge(s)  # We can call a method on the class
    # We can also return instances of the class
    # from TorchScript function/methods
    return s2.clone(), s2.top()

stack, top = do_stacks(torch.classes.my_classes.MyStackClass(["wow"]))
assert top == "wow"
for expected in ["wow", "mom", "hi"]:
    assert stack.pop() == expected

使用自定义类保存、加载和运行 TorchScript 代码

我们还可以在 C++进程中使用自定义注册的 C++类使用 libtorch。例如，让我们定义一个简单的nn.Module，该模块实例化并调用我们的 MyStackClass 类的方法：

import torch

torch.classes.load_library('build/libcustom_class.so')

class Foo(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, s: str) -> str:
        stack = torch.classes.my_classes.MyStackClass(["hi", "mom"])
        return stack.pop() + s

scripted_foo = torch.jit.script(Foo())
print(scripted_foo.graph)

scripted_foo.save('foo.pt')

我们的文件系统中的foo.pt现在包含了我们刚刚定义的序列化的 TorchScript 程序。

现在，我们将定义一个新的 CMake 项目，以展示如何加载这个模型及其所需的.so 文件。有关如何执行此操作的完整说明，请查看在 C++中加载 TorchScript 模型的教程。

与之前类似，让我们创建一个包含以下内容的文件结构：

cpp_inference_example/
  infer.cpp
  CMakeLists.txt
  foo.pt
  build/
  custom_class_project/
    class.cpp
    CMakeLists.txt
    build/

请注意，我们已经复制了序列化的foo.pt文件，以及上面的custom_class_project的源代码树。我们将把custom_class_project作为这个 C++项目的依赖项，以便我们可以将自定义类构建到二进制文件中。

让我们用以下内容填充infer.cpp：

#include  <torch/script.h>

#include  <iostream>
#include  <memory>

int  main(int  argc,  const  char*  argv[])  {
  torch::jit::Module  module;
  try  {
  // Deserialize the ScriptModule from a file using torch::jit::load().
  module  =  torch::jit::load("foo.pt");
  }
  catch  (const  c10::Error&  e)  {
  std::cerr  <<  "error loading the model\n";
  return  -1;
  }

  std::vector<c10::IValue>  inputs  =  {"foobarbaz"};
  auto  output  =  module.forward(inputs).toString();
  std::cout  <<  output->string()  <<  std::endl;
}

类似地，让我们定义我们的 CMakeLists.txt 文件：

cmake_minimum_required(VERSION  3.1  FATAL_ERROR)
project(infer)

find_package(Torch  REQUIRED)

add_subdirectory(custom_class_project)

# Define our library target
add_executable(infer  infer.cpp)
set(CMAKE_CXX_STANDARD  14)
# Link against LibTorch
target_link_libraries(infer  "${TORCH_LIBRARIES}")
# This is where we link in our libcustom_class code, making our
# custom class available in our binary.
target_link_libraries(infer  -Wl,--no-as-needed  custom_class)

您知道该怎么做：cd build，cmake，然后make：

$  cd  build
$  cmake  -DCMAKE_PREFIX_PATH="$(python  -c  'import torch.utils; print(torch.utils.cmake_prefix_path)')"  ..
  --  The  C  compiler  identification  is  GNU  7.3.1
  --  The  CXX  compiler  identification  is  GNU  7.3.1
  --  Check  for  working  C  compiler:  /opt/rh/devtoolset-7/root/usr/bin/cc
  --  Check  for  working  C  compiler:  /opt/rh/devtoolset-7/root/usr/bin/cc  --  works
  --  Detecting  C  compiler  ABI  info
  --  Detecting  C  compiler  ABI  info  -  done
  --  Detecting  C  compile  features
  --  Detecting  C  compile  features  -  done
  --  Check  for  working  CXX  compiler:  /opt/rh/devtoolset-7/root/usr/bin/c++
  --  Check  for  working  CXX  compiler:  /opt/rh/devtoolset-7/root/usr/bin/c++  --  works
  --  Detecting  CXX  compiler  ABI  info
  --  Detecting  CXX  compiler  ABI  info  -  done
  --  Detecting  CXX  compile  features
  --  Detecting  CXX  compile  features  -  done
  --  Looking  for  pthread.h
  --  Looking  for  pthread.h  -  found
  --  Looking  for  pthread_create
  --  Looking  for  pthread_create  -  not  found
  --  Looking  for  pthread_create  in  pthreads
  --  Looking  for  pthread_create  in  pthreads  -  not  found
  --  Looking  for  pthread_create  in  pthread
  --  Looking  for  pthread_create  in  pthread  -  found
  --  Found  Threads:  TRUE
  --  Found  torch:  /local/miniconda3/lib/python3.7/site-packages/torch/lib/libtorch.so
  --  Configuring  done
  --  Generating  done
  --  Build  files  have  been  written  to:  /cpp_inference_example/build
$  make  -j
  Scanning  dependencies  of  target  custom_class
  [  25%]  Building  CXX  object  custom_class_project/CMakeFiles/custom_class.dir/class.cpp.o
  [  50%]  Linking  CXX  shared  library  libcustom_class.so
  [  50%]  Built  target  custom_class
  Scanning  dependencies  of  target  infer
  [  75%]  Building  CXX  object  CMakeFiles/infer.dir/infer.cpp.o
  [100%]  Linking  CXX  executable  infer
  [100%]  Built  target  infer

现在我们可以运行我们令人兴奋的 C++二进制文件了：

$  ./infer
  momfoobarbaz

令人难以置信！

将自定义类移动到/从 IValues

还可能需要将自定义类移入或移出IValue，例如当您从 TorchScript 方法中获取或返回IValue时，或者您想在 C++中实例化自定义类属性时。要从自定义 C++类实例创建IValue：

  • torch::make_custom_class<T>()提供了类似于 c10::intrusive_ptr的 API，它将接受您提供的一组参数，调用与该参数集匹配的 T 的构造函数，并将该实例包装起来并返回。但是，与仅返回自定义类对象的指针不同，它返回包装对象的IValue。然后，您可以直接将此IValue传递给 TorchScript。

  • 如果您已经有一个指向您的类的intrusive_ptr，则可以直接使用构造函数IValue(intrusive_ptr<T>)从中构造一个 IValue。

将IValue转换回自定义类：

  • IValue::toCustomClass<T>()将返回指向IValue包含的自定义类的intrusive_ptr<T>。在内部，此函数正在检查T是否已注册为自定义类，并且IValue确实包含自定义类。您可以通过调用isCustomClass()手动检查IValue是否包含自定义类。

为自定义 C++类定义序列化/反序列化方法

如果尝试将具有自定义绑定的 C++类作为属性保存为ScriptModule，将会收到以下错误：

# export_attr.py
import torch

torch.classes.load_library('build/libcustom_class.so')

class Foo(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.stack = torch.classes.my_classes.MyStackClass(["just", "testing"])

    def forward(self, s: str) -> str:
        return self.stack.pop() + s

scripted_foo = torch.jit.script(Foo())

scripted_foo.save('foo.pt')
loaded = torch.jit.load('foo.pt')

print(loaded.stack.pop()) 

$  python  export_attr.py
RuntimeError:  Cannot  serialize  custom  bound  C++  class  __torch__.torch.classes.my_classes.MyStackClass.  Please  define  serialization  methods  via  def_pickle  for  this  class.  (pushIValueImpl  at  ../torch/csrc/jit/pickler.cpp:128)

这是因为 TorchScript 无法自动确定从您的 C++类中保存哪些信息。您必须手动指定。方法是在类上使用class_的特殊def_pickle方法定义__getstate__和__setstate__方法。

注意

TorchScript 中__getstate__和__setstate__的语义与 Python pickle 模块的相同。您可以关于我们如何使用这些方法。

这里是我们可以添加到MyStackClass注册中的def_pickle调用的示例，以包含序列化方法：

 // class_<>::def_pickle allows you to define the serialization
  // and deserialization methods for your C++ class.
  // Currently, we only support passing stateless lambda functions
  // as arguments to def_pickle
  .def_pickle(
  // __getstate__
  // This function defines what data structure should be produced
  // when we serialize an instance of this class. The function
  // must take a single `self` argument, which is an intrusive_ptr
  // to the instance of the object. The function can return
  // any type that is supported as a return value of the TorchScript
  // custom operator API. In this instance, we've chosen to return
  // a std::vector<std::string> as the salient data to preserve
  // from the class.
  [](const  c10::intrusive_ptr<MyStackClass<std::string>>&  self)
  ->  std::vector<std::string>  {
  return  self->stack_;
  },
  // __setstate__
  // This function defines how to create a new instance of the C++
  // class when we are deserializing. The function must take a
  // single argument of the same type as the return value of
  // `__getstate__`. The function must return an intrusive_ptr
  // to a new instance of the C++ class, initialized however
  // you would like given the serialized state.
  [](std::vector<std::string>  state)
  ->  c10::intrusive_ptr<MyStackClass<std::string>>  {
  // A convenient way to instantiate an object and get an
  // intrusive_ptr to it is via `make_intrusive`. We use
  // that here to allocate an instance of MyStackClass<std::string>
  // and call the single-argument std::vector<std::string>
  // constructor with the serialized state.
  return  c10::make_intrusive<MyStackClass<std::string>>(std::move(state));
  });

注意

我们在 pickle API 中采用了与 pybind11 不同的方法。而 pybind11 有一个特殊函数pybind11::pickle()，您可以将其传递给class_::def()，我们为此目的有一个单独的方法def_pickle。这是因为名称torch::jit::pickle已经被使用，我们不想引起混淆。

一旦以这种方式定义了（反）序列化行为，我们的脚本现在可以成功运行：

$  python  ../export_attr.py
testing

定义接受或返回绑定的 C++类的自定义运算符

一旦定义了自定义 C++类，您还可以将该类用作自定义运算符（即自由函数）的参数或返回值。假设您有以下自由函数：

c10::intrusive_ptr<MyStackClass<std::string>>  manipulate_instance(const  c10::intrusive_ptr<MyStackClass<std::string>>&  instance)  {
  instance->pop();
  return  instance;
}

您可以在TORCH_LIBRARY块内运行以下代码来注册它：

 m.def(
  "manipulate_instance(__torch__.torch.classes.my_classes.MyStackClass x) -> __torch__.torch.classes.my_classes.MyStackClass Y",
  manipulate_instance
  );

有关注册 API 的更多详细信息，请参考自定义操作教程。

完成后，您可以像以下示例一样使用该运算符：

class TryCustomOp(torch.nn.Module):
    def __init__(self):
        super(TryCustomOp, self).__init__()
        self.f = torch.classes.my_classes.MyStackClass(["foo", "bar"])

    def forward(self):
        return torch.ops.my_classes.manipulate_instance(self.f)

注意

接受 C++类作为参数的运算符的注册要求自定义类已经注册。您可以通过确保自定义类注册和您的自由函数定义位于同一个TORCH_LIBRARY块中，并且自定义类注册位于首位来强制执行此要求。在未来，我们可能会放宽此要求，以便可以以任何顺序注册这些内容。

结论

本教程向您展示了如何将一个 C++类暴露给 TorchScript（以及 Python），如何注册其方法，如何从 Python 和 TorchScript 中使用该类，以及如何使用该类保存和加载代码，并在独立的 C++进程中运行该代码。现在，您可以准备使用与第三方 C++库进行交互的 C++类来扩展您的 TorchScript 模型，或者实现任何其他需要在 Python、TorchScript 和 C++之间平滑过渡的用例。

如常，如果您遇到任何问题或有疑问，您可以使用我们的论坛或GitHub 问题页面联系我们。此外，我们的常见问题（FAQ）页面可能包含有用的信息。


http://www.kler.cn/a/229866.html

相关文章：

  • 05_JavaScript注释与常见输出方式
  • 【C++进阶篇】像传承家族宝藏一样理解C++继承
  • 双十一线上服务调用链路追踪SkyWalking实战分析
  • 第 37 章 - Go 语言 持续集成与持续部署
  • BUUCTF—Reverse—不一样的flag（7）
  • VS的安装和配置
  • vue3项目中使用mapv
  • 基于QPSO-LSTM的短期风电负荷MATLAB预测程序
  • 服务器和CDN推荐
  • http+域名+端口
  • 优雅的从ＨuggingFace下载模型
  • React18构建Vite+Electron项目以及打包
  • 假期算法提升（带你彻底掌握滑动窗口）
  • CSS：水平垂直居中
  • 【Jenkins】pipeline基本使用
  • Java赋能：大学生成绩量化新篇章
  • [英语学习][27][Word Power Made Easy]的精读与翻译优化
  • 第5节、S曲线加减速转动【51单片机+L298N步进电机系列教程】
  • 如何把vue项目打包成桌面程序 electron-builder
  • Windows中如何使用 Anaconda 和 gempy地质建模
  • 中移在线：基于openGauss的数据库自主创新替代实践
  • 【React】前端React 代码中预览展示excel文件
  • RabbitMQ：分布式系统中的高效消息队列
  • 19：Web开发模式与MVC设计模式-Java Web
  • 第三百一十一回
  • 前端工程化之：webpack1-11(其他配置)