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

原文： PyTorch 使用自定义 C ++运算符扩展 TorchScript

PyTorch 1.0 版本向 PyTorch 引入了一种新的编程模型，称为 TorchScript 。 TorchScript 是 Python 编程语言的子集，可以通过 TorchScript 编译器进行解析，编译和优化。 此外，已编译的 TorchScript 模型可以选择序列化为磁盘文件格式，然后可以从纯 C ++(以及 Python）加载并运行该文件格式以进行推理。

TorchScript 支持torch包提供的大量操作子集，使您可以纯粹表示为 PyTorch 的“标准库”中的一系列张量操作来表示多种复杂模型。 但是，有时您可能需要使用自定义 C ++或 CUDA 函数扩展 TorchScript。 虽然我们建议您仅在无法(简单有效地）将您的想法表达为简单的 Python 函数时才诉诸该选项，但我们确实提供了一个非常友好且简单的界面，用于使用 ATen 定义自定义 C ++和 CUDA 内核。 ，PyTorch 的高性能 C ++张量库。 绑定到 TorchScript 后，您可以将这些自定义内核(或“ ops”）嵌入到 TorchScript 模型中，并以 Python 或直接以 C ++的序列化形式执行它们。

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

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

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

#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=*/{8, 8});


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

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

小费

有关 ATen 的更多信息，请参见本说明，ATen 是为 PyTorch 提供Tensor类的库。 此外，本教程的描述了如何在 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<float>());

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

cv::Mat warp_mat(/*rows=*/warp.size(0),
                 /*cols=*/warp.size(1),
                 /*type=*/CV_32FC1,
                 /*data=*/warp.data<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::from_blob(output_mat.ptr<float>(), /*sizes=*/{8, 8})

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

从我们的运算符实现返回该张量之前，我们必须在张量上调用.clone()以执行基础数据的存储副本。 这样做的原因是torch::from_blob返回的张量不拥有其数据。 那时，数据仍归 OpenCV 矩阵所有。 但是，此 OpenCV 矩阵将超出范围，并在函数末尾重新分配。 如果我们按原样返回output张量，那么当我们在函数外使用它时，它将指向无效的内存。 调用.clone()将返回一个新的张量，其中包含新张量自己拥有的原始数据的副本。 因此，返回外部世界是安全的。

使用 TorchScript 注册自定义运算符

现在，已经在 C ++中实现了自定义运算符，我们需要在 TorchScript 运行时和编译器中将注册为。 这将使 TorchScript 编译器可以在 TorchScript 代码中解析对我们自定义运算符的引用。 注册非常简单。 对于我们的情况，我们需要编写：

static auto registry =
  torch::RegisterOperators("my_ops::warp_perspective", &warp_perspective);

op.cpp文件的全局范围内的某个位置。 这将创建一个全局变量registry，该变量将在其构造函数中向 TorchScript 注册我们的运算符(即每个程序一次）。 我们指定运算符的名称，以及指向其实现的指针(我们之前编写的函数）。 该名称包括两部分：命名空间(my_ops）和我们正在注册的特定运算符的名称(warp_perspective）。 名称空间和操作员名称由两个冒号(::）分隔。

Tip

如果要注册多个运算符，可以在构造函数之后将调用链接到.op()：

static auto registry =
  torch::RegisterOperators("my_ops::warp_perspective", &warp_perspective)
  .op("my_ops::another_op", &another_op)
  .op("my_ops::and_another_op", &and_another_op);

在后台，RegisterOperators将执行许多相当复杂的 C ++模板元编程魔术技巧，以推断我们传递给它的函数指针的参数和返回值类型(&warp_perspective）。 此信息用于为我们的操作员形成功能模式。 函数模式是操作员的结构化表示形式，一种“签名”或“原型”，由 TorchScript 编译器用来验证 TorchScript 程序的正确性。

建立自定义操作员

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

用 CMake 构建

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

warp-perspective/
  op.cpp
  CMakeLists.txt

另外，请确保从 pytorch.org 中获取 LibTorch 发行版的最新版本，该软件包打包了 PyTorch 的 C ++库和 CMake 构建文件。 将解压缩的发行版放置在文件系统中可访问的位置。 以下段落将将该位置称为/path/to/libtorch。 我们的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++11
target_compile_features(warp_perspective PRIVATE cxx_range_for)
## Link against LibTorch
target_link_libraries(warp_perspective "${TORCH_LIBRARIES}")
## Link against OpenCV
target_link_libraries(warp_perspective opencv_core opencv_imgproc)

警告

此设置对构建环境进行了一些假设，特别是有关 OpenCV 安装的假设。 上面的CMakeLists.txt文件已在运行 Ubuntu Xenial 的 Docker 容器中通过apt安装了libopencv-dev进行了测试。 如果它对您不起作用，并且您感到困惑，请使用随附的教程资料库中的Dockerfile构建一个隔离的，可复制的环境，在其中可以使用本教程中的代码。 如果您遇到其他麻烦，请在教程资料库中提交问题，或在我们的论坛中发布问题。

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

$ mkdir build
$ cd build
$ cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
-- 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命令中，应将/path/to/libtorch替换为未压缩的 LibTorch 发行版的路径。

我们将在下面进一步探讨如何使用和调用我们的运算符，但是为了早日获得成功，我们可以尝试在 Python 中运行以下代码：

>>> import torch
>>> torch.ops.load_library("/path/to/libwarp_perspective.so")
>>> print(torch.ops.my_ops.warp_perspective)

在这里，/path/to/libwarp_perspective.so应该是我们刚刚构建的libwarp_perspective.so共享库的相对或绝对路径。 如果一切顺利，这应该打印类似

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

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

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

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

您已经了解了如何将运算符导入 Python：torch.ops.load_library()。 此函数采用包含自定义运算符的共享库的路径，并将其加载到当前进程中。 加载共享库还将执行我们放入自定义运算符实现文件中的全局RegisterOperators对象的构造函数。 这将在 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("libwarp_perspective.so")
>>> 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 ops 不支持它。

从这里开始，您可以在脚本或跟踪代码中使用自定义运算符，就像torch包中的其他函数一样。 实际上，诸如torch.matmul之类的“标准库”功能与自定义运算符的注册路径大致相同，这使得自定义运算符在 TorchScript 中的使用方式和位置方面真正成为一等公民。

使用自定义运算符进行跟踪

首先，将我们的运算符嵌入到跟踪函数中。 回想一下，为了进行跟踪，我们从一些原始的 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)
    %y : Float(8, 5)
    %z : Float(4, 5)) {
  %3 : Float(4, 5) = aten::matmul(%x, %y)
  %4 : Float(4, 5) = aten::relu(%z)
  %5 : int = prim::Constant[value=1]()
  %6 : Float(4, 5) = aten::add(%3, %4, %5)
  return (%6);
}

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

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


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)
    %y : Float(8, 5)
    %z : Float(8, 5)) {
    %3 : int = prim::Constant[value=3]()
    %4 : int = prim::Constant[value=6]()
    %5 : int = prim::Constant[value=0]()
    %6 : int[] = prim::Constant[value=[0, -1]]()
    %7 : Float(3, 3) = aten::eye(%3, %4, %5, %6)
    %x : Float(8, 8) = my_ops::warp_perspective(%x.1, %7)
    %11 : Float(8, 5) = aten::matmul(%x, %y)
    %12 : Float(8, 5) = aten::relu(%z)
    %13 : int = prim::Constant[value=1]()
    %14 : Float(8, 5) = aten::add(%11, %12, %13)
    return (%14);
  }

如此简单地将 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 ++中的RegisterOperators对象注册的实现，并将其编译为图形表示形式：

>>> 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的引用。

Attention

TorchScript 图形表示仍可能更改。 不要依靠它看起来像这样。

在 Python 中使用自定义运算符时，确实如此。 简而言之，您可以使用torch.ops.load_library导入包含运算符的库，并像其他任何torch运算符一样，从跟踪或编写脚本的 TorchScript 代码中调用自定义操作。

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

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

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

Note

从技术上讲，您还可以在运行时将共享库动态加载到 C ++应用程序中，就像在 Python 中一样。 在 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().
  std::shared_ptr<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=/path/to/libtorch ..
-- 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共享库链接起来。

Attention

上面的示例中嵌入了一个关键细节：warp_perspective链接行的-Wl,--no-as-needed前缀。 这是必需的，因为我们实际上不会在应用程序代码中从warp_perspective共享库中调用任何函数。 我们只需要运行全局RegisterOperators对象的构造函数即可。 麻烦的是，这使链接器感到困惑，并使其认为可以完全跳过针对库的链接。 在 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=/path/to/libtorch ..
-- 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 ++库进行接口的 C ++运算符扩展 TorchScript 模型，编写自定义的高性能 CUDA 内核，或实现任何其他需要 Python，TorchScript 和 C ++之间的界线才能平稳融合的用例。

与往常一样，如果您遇到任何问题或疑问，可以使用我们的论坛或 GitHub 问题进行联系。 另外，我们的常见问题解答(FAQ）页面可能包含有用的信息。

附录 A：建立自定义操作员的更多方法

“构建自定义运算符”一节介绍了如何使用 CMake 将自定义运算符构建到共享库中。 本附录概述了两种进一步的编译方法。 他们俩都使用 Python 作为编译过程的“驱动程序”或“接口”。 此外，两者都重新使用了现有基础结构 PyTorch 提供了 C ++扩展 ，它们是依赖于 [pybind11 用于将功能从 C ++“显式”绑定到 Python。

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

使用 JIT 编译进行构建

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

Note

这里的“ JIT 编译”与 TorchScript 编译器中用于优化程序的 JIT 编译无关。 这只是意味着您的自定义运算符 C ++代码将在您首次导入时在系统 <cite>/ tmp</cite> 目录下的文件夹中编译，就像您自己事先对其进行编译一样。

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


static auto registry =
  torch::RegisterOperators("my_ops::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 模块而不是普通的共享库(它们没有 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.custom.warp_perspective)
<built-in method custom::warp_perspective of PyCapsule object at 0x7ff51c5b7bd0>