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

创建时间： 2018-11-28 |上次更新时间： 2024 年 7 月 22 日 |上次验证： Nov 05， 2024

警告

本教程从 PyTorch 2.4 开始已弃用。请参阅 PyTorch 自定义运算符，了解有关 PyTorch 自定义运算符的最新指南。

PyTorch 1.0 版本向 PyTorch 引入了一个名为 TorchScript 的新编程模型。TorchScript 是一个 Python 编程语言的子集，可以解析、编译和 由 TorchScript 编译器优化。此外，编译后的 TorchScript 模型具有 序列化为磁盘文件格式的选项，您可以 随后从纯 C++（以及 Python）加载和运行以进行推理。

TorchScript 支持包提供的大量操作，允许您将多种复杂模型纯粹表示为一个系列 来自 PyTorch 的“标准库”的张量操作。尽管如此，可能会 有时你发现自己需要用自定义的 C++ 或 CUDA 函数。虽然我们建议您仅在以下情况下使用此选项 你的想法不能（足够有效地）表达为一个简单的 Python 函数， 我们确实提供了一个非常友好和简单的接口来定义自定义 C++ 和 使用 ATen 的 CUDA 内核，PyTorch 的高 performance C++ 张量库。绑定到 TorchScript 后，您可以将这些 自定义内核（或“操作”）添加到您的 TorchScript 模型中，并在 Python 及其序列化形式直接在 C++ 中。torch

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

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

在本教程中，我们将公开 warpPerspective 函数，该函数将透视变换应用于图像，从 OpenCV 到 TorchScript 作为自定义运算符。第一步是编写 implementation 我们在 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 头文件 ，以及头文件，它公开了 PyTorch 的 C++ API 中的所有必要功能，我们 需要编写自定义 TorchScript 运算符。我们的函数有两个参数：一个 input 和 transformation matrix 我们希望应用于图像。这些输入的类型为 ， PyTorch 在 C++ 中的张量类型（这也是所有张量的基础类型 在 Python 中）。我们函数的返回类型也将是一个 .opencv2/opencv.hpptorch/script.hwarp_perspectiveimagewarptorch::Tensorwarp_perspectivetorch::Tensor

提示

请参阅此说明 有关 ATen 的更多信息，该库将类提供给 PyTorch 的 Torch 中。此外，本教程还介绍了如何 在 C++ 中分配和初始化新的 Tensor 对象（不需要） 运算符）。Tensor

注意力

TorchScript 编译器可以理解固定数量的类型。仅这些类型 可用作自定义运算符的参数。目前，这些类型是：、、 和 s 这些类型。请注意，仅支持 only 和 not 、 以及 only 和 not 其他整型类型，例如 、 或 。torch::Tensortorch::Scalardoubleint64_tstd::vectordoublefloatint64_tintshortlong

在我们的函数中，我们需要做的第一件事是将 PyTorch tensors 转换为 OpenCV 矩阵，因为 OpenCV 期望对象作为输入。幸运的是，有一种方法可以在不复制的情况下做到这一点 任何数据。在前几行中，warpPerspectivecv::Mat

  cv::Mat image_mat(/*rows=*/image.size(0),
                    /*cols=*/image.size(1),
                    /*type=*/CV_32FC1,
                    /*data=*/image.data_ptr<float>());

我们调用 OpenCV 类的这个构造函数来将我们的张量转换为对象。我们通过 它是原始张量的行数和列数，数据类型 （我们将像本例一样修复），最后是指向 基础数据 – a .这个构造函数有什么特别之处 类是它不会复制 Importing 数据。相反，它会 只需引用此内存即可对 执行的所有操作进行 .如果 就地操作，这将反映在 原始张量（反之亦然）。这允许我们调用 后续 OpenCV 例程替换为库的原生矩阵类型，即使 我们实际上是将数据存储在 PyTorch 张量中。我们重复此过程以 将 PyTorch 张量转换为 OpenCV 矩阵：MatMatimagefloat32float*MatMatimage_matimagewarpwarp_mat

  cv::Mat warp_mat(/*rows=*/warp.size(0),
                   /*cols=*/warp.size(1),
                   /*type=*/CV_32FC1,
                   /*data=*/warp.data_ptr<float>());

接下来，我们准备调用我们渴望使用的 OpenCV 函数 TorchScript： .为此，我们将 OpenCV 函数 the 和 矩阵以及一个空的输出矩阵 叫。我们还指定了我们想要输出的大小 matrix （image） 来。对于此示例，它被硬编码为：warpPerspectiveimage_matwarp_matoutput_matdsize8 x 8

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

自定义算子实现的最后一步是将 back 转换为 PyTorch 张量，以便我们可以在 PyTorch 的 Torch 中。这与我们之前在 另一个方向。在这种情况下，PyTorch 提供了一个方法。在这种情况下，blob 表示指向内存的某个不透明的平面指针，该 我们希望解释为 PyTorch 张量。对 looks 的调用 喜欢这个：output_mattorch::from_blobtorch::from_blob

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

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

在从我们的算子实现返回这个张量之前，我们必须调用这个张量来执行底层数据的内存复制。这 这样做的原因是，返回一个不拥有 它的数据。此时，数据仍归 OpenCV 矩阵所有。然而 这个 OpenCV 矩阵将超出范围，并在 功能。如果我们按原样返回张量，它将指向无效 memory 的调用。调用返回 具有新张量拥有自身的原始数据副本的新张量。 因此，返回外面的世界是安全的。.clone()torch::from_bloboutput.clone()

使用 TorchScript 注册自定义运算符

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

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

在我们文件的顶层的某个位置。宏 创建一个函数，该函数将在程序启动时调用。名称 的 （） 作为第一个参数给出（它不应该 放在引号中）。第二个参数 （） 定义一个类型的变量，该变量是注册运算符的主接口。 该方法实际上创建了一个名为 ， 将其公开给 Python 和 TorchScript。您可以定义任意数量的运算符 根据需要对 .op.cppTORCH_LIBRARYmy_opsmtorch::LibraryLibrary::defwarp_perspectivedef

在幕后，该函数实际上正在做相当多的工作： 它使用模板元编程来检查您的 函数并将其转换为指定运算符 type 中。def

构建自定义运算符

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

环境设置

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

conda install -c pytorch pytorch
conda install opencv

使用 CMake 构建

要使用 CMake 构建系统将我们的自定义运算符构建到共享库中，我们需要编写一个短文件并将其与上一个文件放在一起。为此，让我们就 如下所示的目录结构：CMakeLists.txtop.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)

现在要构建我们的 Operator，我们可以从我们的文件夹运行以下命令：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

，这会将共享库文件放置在文件夹中。在上面的命令中，我们使用 helper 变量来方便地告诉我们在哪里 用于 PyTorch 安装的 cmake 文件是。libwarp_perspective.sobuildcmaketorch.utils.cmake_prefix_path

我们将在下面进一步详细探讨如何使用和调用我们的 operator，但要 获得早期的成功感觉，我们可以尝试在 蟒：

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>

这是我们稍后将用于调用 custom 运算符的 Python 函数。

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

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

您已经了解了如何将运算符导入 Python： 。此函数采用共享库的路径 包含自定义运算符，并将其加载到当前进程中。加载 shared library 也将执行该块。这将注册 我们的自定义运算符与 TorchScript 编译器一起使用，并允许我们使用它 运算符。torch.ops.load_library()TORCH_LIBRARY

您可以将 loaded 运算符称为 ， 其中 是 Operator Name 的 namespace 部分，以及 Operator 的 function name。对于运算符，我们编写了 上面，命名空间为 WAS 和函数名称 ， 这意味着我们的运算符可以作为 . 虽然此函数可用于脚本化或跟踪的 TorchScript 模块，但我们 也可以在原版 Eager PyTorch 中使用它，并传递常规 PyTorch 张：torch.ops.<namespace>.<function><namespace><function>my_opswarp_perspectivetorch.ops.my_ops.warp_perspective

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 中访问 TorchScript 编译器（在 C++ 中 land） 将查看函数是否已注册，而 如果是这样，则返回此函数的 Python 句柄，我们随后可以使用该句柄 从 Python 调用我们的 C++ 运算符实现。这是一个值得注意的 TorchScript 自定义运算符和 C++ 扩展之间的区别：C++ 扩展是使用 Pybind11 手动绑定的，而 TorchScript 自定义操作是 由 PyTorch 本身动态绑定。Pybind11 为您提供了更大的灵活性 关于您可以绑定到 Python 中的类型和类，因此 建议用于纯 Eager 代码，但 TorchScript 不支持 老年 退休金 计划。torch.ops.namespace.functionnamespace::function

从这里开始，您可以在脚本化或跟踪代码中使用自定义运算符 就像你对包中的其他函数一样。事实上，“标准 library“ 函数与 go 一样 注册路径作为自定义运算符，这使得自定义运算符真正 一流的公民，在如何以及在哪里使用它们 TorchScript 的 TorchScript 中。（但是，一个区别是标准库函数 具有自定义编写的 Python 参数解析逻辑，该逻辑与参数解析不同。torchtorch.matmultorch.ops

将 Custom 运算符与跟踪一起使用

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

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

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

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.relutorch

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 代码中就这么简单！

将 Custom Operator 与脚本结合使用

除了跟踪之外，还有另一种方法可以得到 PyTorch 程序是直接在 TorchScript 中编写你的代码。TorchScript 是 主要是 Python 语言的一个子集，但有一些限制使其 更容易让 TorchScript 编译器推理程序。你把你的 常规 PyTorch 代码添加到 TorchScript 中，方法是使用 for free 函数和 for 类中的方法（也必须派生自 ）。有关的更多详细信息，请参阅此处 TorchScript 注解。@torch.jit.script@torch.jit.script_methodtorch.jit.ScriptModule

使用 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

这会及时将函数编译成一个图形 表示，我们可以在 property 中检查它：computecompute.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 编译器看到对 的引用时，它会找到 通过 C++ 中的函数注册，并将其编译成 图表表示：torch.ops.my_ops.warp_perspectiveTORCH_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 中使用我们的自定义运算符时，这就是真正的意义所在。在 short，您可以使用 导入包含运算符的库，并像从跟踪或脚本化的 TorchScript 代码中调用任何其他运算符一样调用自定义运算。torch.ops.load_librarytorch

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

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

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

注意

从技术上讲，您还可以将共享库动态加载到 C++ 中 应用程序，这与我们在 Python 中所做的方式大致相同。在 Linux 上，您可以使用 dlopen 执行此操作。存在 其他平台上的等效项。

基于上面链接的 C++ 执行教程，让我们从最小的 C++ 应用程序放在一个文件中，与我们的 custom 运算符，该运算符加载并执行序列化的 TorchScript 模型：main.cpp

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

或者也许不是。也许现在还不行。答案是肯定的！我们尚未链接自定义 Operator 库。让我们现在就做这件事，并且要这样做 适当地让我们稍微更新一下我们的文件组织，如下所示：

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

这将允许我们将库 CMake 目标添加为 子目录中。文件夹中的顶层应如下所示：warp_perspectiveCMakeLists.txtexample_app

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 配置看起来和以前很像，只是我们将 CMake 构建添加为子目录。一旦它的 CMake 代码运行，我们 将我们的应用程序与共享的 库。warp_perspectiveexample_appwarp_perspective

注意力

上面的例子中嵌入了一个关键细节：链接行的前缀。这是 required，因为我们实际上不会在应用程序代码中从共享库中调用任何函数。我们只需要运行函数。不便的是，这个 混淆链接器，并使其认为它可以跳过对 库。在 Linux 上，标志会强制链接 发生（注意：此标志特定于 Linux！还有其他解决方法 为了这个。最简单的是在 operator 库中定义一些函数 ，您需要从主应用程序调用。这可以像 function 声明在某个 Headers 中声明，然后定义在 Operator 库中。调用此函数 在主应用程序中，链接器会给链接器留下这样的印象，即这是一个 值得链接的库。不幸的是，这超出了我们的控制范围。 我们宁愿让您知道原因和简单的解决方法 而不是给你一些不透明的宏来 plop 在你的代码中。-Wl,--no-as-neededwarp_perspectivewarp_perspectiveTORCH_LIBRARY-Wl,--no-as-neededvoid init();void init() { }init()

现在，由于我们现在在顶层找到了包，因此子目录中的文件可以是 缩短了一点。它应该看起来像这样：TorchCMakeLists.txtwarp_perspective

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++，如何将其构建到共享库中，如何在 Python 中使用它来定义 TorchScript 模型，最后如何将其加载到 C++ 应用程序中 推理工作负载。现在，您可以使用 与第三方 C++ 库接口的 C++ 运算符，编写自定义高 performance CUDA 内核，或实现任何其他需要 在 Python、TorchScript 和 C++ 之间平滑混合。

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

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

“构建自定义 Operator” 一节介绍了如何构建自定义 运算符添加到使用 CMake 的共享库中。本附录进一步概述了两个 方法进行编译。他们都使用 Python 作为 “驱动程序” 或 “interface” 添加到编译过程。此外，两者都重用了现有的 基础设施 PyTorch 提供 *C++ 扩展*，它们是 vanilla （eager） PyTorch 相当于 TorchScript 自定义运算符，依赖于 pybind11 进行 “显式” 绑定 函数从 C++ 转换为 Python。

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

使用 JIT 编译进行构建

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

注意

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

此 JIT 编译功能有两种类型。在第一个 API 中，你仍然保持 你的 operator 实现放在一个单独的文件 （） 中，然后用于编译你的扩展。通常，这个 function 将返回 Python 模块，公开您的 C++ 扩展。然而 由于我们没有将自定义运算符编译成自己的 Python 模块，因此 只想编译一个普通的共享库。幸运的是，有一个参数 我们可以设置为 to 来表示我们只对构建 共享库，而不是 Python 模块。 然后将编译共享库并将其加载到当前进程中， 就像以前一样：op.cpptorch.utils.cpp_extension.load()torch.utils.cpp_extension.load()is_python_moduleFalsetorch.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 中使用它，则不应执行 多次注册的 Cell，因为每次执行都注册 new 库并重新注册 custom 运算符。如果需要重新执行它， 请事先重新启动笔记本的 Python 内核。

使用 Setuptools 进行构建

仅从 Python 构建自定义运算符的第二种方法是 以使用 .这样做的好处是 用于构建用 C++ 编写的 Python 模块的强大而广泛的接口。 但是，由于 this 实际上是用于构建 Python 模块和 不是普通的共享库（没有必要的入口点 Python 期望从模块中），这条路由可能会有点古怪。也就是说，你 need 是一个文件来代替它，它看起来像 这：setuptoolssetuptoolssetuptoolssetup.pyCMakeLists.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)},
)

请注意，我们在底部的 中启用了该选项。这指示省略任何 生成的共享库名称中的 Python-3 特定 ABI 后缀。 否则，例如在 Python 3.7 上，该库可能被称为 ABI 标签在哪里，但我们真的只希望它 名叫no_python_abi_suffixBuildExtensionsetuptoolswarp_perspective.cpython-37m-x86_64-linux-gnu.socpython-37m-x86_64-linux-gnuwarp_perspective.so

如果我们现在从 文件夹中，我们应该看到如下内容：python setup.py build developsetup.py

$ 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

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

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