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

使用自定义 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 后，您可以将这些自定义内核（或“操作”）嵌入到 TorchScript 模型中，并以 Python 或直接以 C++ 的序列化形式执行它们。

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

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

在本教程中，我们将公开warpPerspective函数，该函数将透视转换应用于图像，从 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，这是 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_ptr<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_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()方法来获取指向基础数据的原始指针（就像之前的 PyTorch 张量的.data_ptr()一样）。 我们还指定了张量的输出形状，我们将其硬编码为8 x 8。 然后torch::from_blob的输出是torch::Tensor，指向 OpenCV 矩阵拥有的内存。

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

使用 TorchScript 注册自定义运算符

现在，已经在 C++ 中实现了自定义运算符，我们需要在 T​​orchScript 运行时和编译器中将其注册。 这将使 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 或setuptools之类的 Python 替代方法来构建我们的运算符。 为简洁起见，以下段落仅讨论 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

这会将libwarp_perspective.so共享库文件放置在build文件夹中。 在上面的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..，其中是运算符名称的名称空间部分，而是运算符的函数名称。 对于我们上面编写的运算符，名称空间为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导入包含您的运算符的库，并像其他任何torch运算符一样，从跟踪或编写脚本的 TorchScript 代码中调用自定义操作。

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

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

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

注意

从技术上讲，您还可以在运行时将共享库动态加载到 C++ 应用中，就像在 Python 中一样。 在 Linux 上，可以使用dlopen来执行此操作。 在其他平台上也存在等效项。

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

#include  // One-stop header.

#include 
#include 

int main(int argc, const char* argv[]) {
  if (argc != 2) {
    std::cerr << "usage: example-app \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)

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

并在尚未通过模型的情况下运行它：

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

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/
  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++ 库进行接口的 C++ 运算符扩展 TorchScript 模型，编写自定义的高性能 CUDA 内核，或实现任何其他需要 Python，TorchScript 和 C++ 之间的界线才能平稳融合的用例。