使用自定义 C++ 操作扩展 TorchScript

创建日期: 2018年11月28日 | 最后更新日期: 2024年7月22日 | 最后验证日期: 2024年11月5日

警告

本教程自 PyTorch 2.4 起已弃用。请参阅 PyTorch 自定义操作，以获取关于 PyTorch 自定义操作的最新指南。

PyTorch 1.0 版本的发布引入了 PyTorch 的一种新的编程模型，称为 TorchScript。TorchScript 是 Python 编程语言的一个子集， 可以通过 TorchScript 编译器进行解析、编译和优化。此外，编译后的 TorchScript 模型可以选择性地序列化为磁盘上的文件格式，您可以随后从纯 C++（以及 Python）中加载并运行这些模型以进行推理。

TorchScript 支持 torch 包提供的大量操作，允许你通过 PyTorch 的“标准库”中的张量操作序列来表达各种复杂的模型。然而，在某些情况下，你可能需要使用自定义的 C++ 或 CUDA 函数扩展 TorchScript。虽然我们建议只有在你的想法无法以简单 Python 函数的形式（高效地）表达时才选择此选项，但我们确实提供了一个非常友好且简单的接口，用于使用 ATen 定义自定义的 C++ 和 CUDA 内核，ATen 是 PyTorch 的高性能 C++ 张量库。一旦绑定到 TorchScript，你可以将这些自定义内核（或“操作”）嵌入到你的 TorchScript 模型中，并在 Python 中以及直接在 C++ 中以序列化形式执行它们。

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

在C++中实现自定义操作

在这个教程中，我们将从 OpenCV 中公开 warpPerspective 函数， 该函数对图像应用透视变换，并将其作为自定义操作符暴露给 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 的信息，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 张量的行数和列数、数据类型（在这个例子中我们将固定为 14”）以及指向底层数据的原始指针（一个 15”）传递给它。16” 类的这个构造函数特别之处在于它不会复制输入数据。相反，它将在对 17” 执行的所有操作中引用这块内存。如果对 18” 进行了原地操作，这将反映在原始 19” 张量上（反之亦然）。这使我们能够使用库的原生矩阵类型调用后续的OpenCV例程，尽管我们实际上是在PyTorch张量中存储数据。我们重复此过程将 20” PyTorch张量转换为 21” OpenCV矩阵：

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

接下来，我们准备调用我们在TorchScript中非常期待使用的OpenCV函数：warpPerspective。为此，我们将OpenCV函数传递给image_mat和warp_mat矩阵，以及一个名为output_mat的空输出矩阵。我们还指定了输出矩阵（图像）的大小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();

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

在从我们的操作符实现中返回此张量之前，我们必须对该张量调用 .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 中进行研究和实验，或者在无 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

这将把一个 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.<namespace>.<function>， 其中 <namespace> 是运算符名称的命名空间部分，而 <function> 是运算符的函数名称。对于上面编写的运算符， 命名空间是 my_ops，函数名称是 warp_perspective， 这意味着我们的运算符可以作为 torch.ops.my_ops.warp_perspective 使用。 虽然这个函数可以在脚本化或追踪的 TorchScript 模块中使用，我们 也可以直接在普通的 eager 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)))

producing:

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 方面提供了更大的灵活性，因此对于纯急切代码（eager code）推荐使用，但它不支持 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)

Producing:

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)

Producing:

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

使用自定义操作符与Script

除了追踪之外，另一种获得PyTorch程序的TorchScript表示的方法是直接用TorchScript编写代码。TorchScript在很大程度上是Python语言的一个子集，有一些限制使得TorchScript编译器更容易理解程序。你可以通过为普通PyTorch代码添加注释来将其转换为TorchScript：对于自由函数使用@torch.jit.script，对于类中的方法（该类还必须从torch.jit.ScriptModule派生）使用@torch.jit.script_method。有关TorchScript注释的更多详情，请参见这里。

使用 TorchScript 而不是 tracing 的一个特别原因是，tracing 无法捕获 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模型的教程，接下来的几段内容将基于此。

简而言之，自定义操作符可以像普通的 torch 操作符一样执行， 即使是从文件中反序列化并在 C++ 中运行。实现这一点的唯一要求是将之前构建的自定义操作符共享库链接到执行模型的 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 目标作为我们应用程序目标的子目录添加进去。顶级 CMakeLists.txt 在 example_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 配置与之前类似，只是我们添加了 warp_perspective CMake 构建作为子目录。一旦其 CMake 代码运行，我们 将我们的 example_app 应用程序与 warp_perspective 共享库链接起来。

注意力

在上面的示例中有一个关键细节：链接行中的 -Wl,--no-as-needed 前缀。这是必需的，因为我们实际上不会在应用程序代码中调用任何来自 warp_perspective 共享库的函数。我们只需要运行 warp_perspective 函数。不幸的是，这会让链接器误以为它可以完全跳过链接该库。在 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++之间平滑融合的用例。

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

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

“构建自定义操作符”部分解释了如何使用CMake将自定义操作符构建到共享库中。本附录介绍了两种进一步的编译方法。这两种方法都使用Python作为编译过程的“驱动程序”或“接口”。此外，两者都重新利用了PyTorch提供的现有基础设施，用于*C++扩展*，这是依赖于pybind11显式绑定C++函数到Python的TorchScript自定义操作符的原生（急切模式）等效实现。

第一种方法使用 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)

自然，如果源代码 reasonably 短的话，最好只使用 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.my_ops.warp_perspective)
<built-in method custom::warp_perspective of PyCapsule object at 0x7ff51c5b7bd0>