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

创建于：2018 年 11 月 28 日 | 最后更新：2024 年 7 月 22 日 | 最后验证：2024 年 11 月 05 日

警告

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

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

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

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

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

在本教程中，我们将从 OpenCV 向 TorchScript 公开 warpPerspective 函数，该函数将透视变换应用于图像，作为自定义运算符。第一步是用 C++ 编写自定义运算符的实现。让我们将此实现的 file 命名为 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（为 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 张量中（反之亦然）。这允许我们使用库的本机矩阵类型调用后续的 OpenCV 例程，即使我们实际上将数据存储在 PyTorch 张量中。我们重复此过程以将 warp PyTorch 张量转换为 warp_mat OpenCV 矩阵

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

接下来，我们准备好调用我们非常渴望在 TorchScript 中使用的 OpenCV 函数：warpPerspective。为此，我们将 image_mat 和 warp_mat 矩阵以及一个名为 output_mat 的空输出矩阵传递给 OpenCV 函数。我们还指定我们希望输出矩阵（图像）的大小 dsize。在本示例中，它被硬编码为 8 x 8

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

自定义运算符实现的最后一步是将 output_mat 转换回 PyTorch 张量，以便我们可以在 PyTorch 中进一步使用它。这与我们之前在另一个方向上进行转换的操作惊人地相似。在这种情况下，PyTorch 提供了一个 torch::from_blob 方法。在这种情况下，blob 旨在表示一些不透明的、指向内存的平面指针，我们希望将其解释为 PyTorch 张量。对 torch::from_blob 的调用如下所示

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

我们使用 OpenCV Mat 类上的 .ptr<float>() 方法来获取指向底层数据的原始指针（就像之前 PyTorch 张量的 .data_ptr<float>() 一样）。我们还指定了张量的输出形状，我们将其硬编码为 8 x 8。torch::from_blob 的输出是 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 中进行研究和实验，或者加载到 C++ 中以在无 Python 环境中进行推理。有多种方法可以使用纯 CMake 或 Python 替代方案（如 setuptools）来构建我们的运算符。为简洁起见，以下段落仅讨论 CMake 方法。本教程的附录深入探讨了其他替代方案。

环境设置

我们需要安装 PyTorch 和 OpenCV。获取这两者的最简单且与平台无关的方式是通过 Conda

conda install -c pytorch pytorch
conda install opencv

使用 CMake 构建

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

warp-perspective/
  op.cpp
  CMakeLists.txt

然后我们的 CMakeLists.txt 文件的内容应如下所示

cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
project(warp_perspective)

find_package(Torch REQUIRED)
find_package(OpenCV REQUIRED)

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

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

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

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

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

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

如果一切顺利，这将打印出类似以下内容

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

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

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

一旦我们的自定义运算符被构建到共享库中，我们就准备好在 Python 中的 TorchScript 模型中使用此运算符。这分为两个部分：首先将运算符加载到 Python 中，其次在 TorchScript 代码中使用该运算符。

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

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

产生

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 自定义 ops 是由 PyTorch 本身动态绑定的。 Pybind11 在您可以绑定到 Python 中的类型和类方面提供了更大的灵活性，因此建议用于纯 eager 代码，但 TorchScript ops 不支持它。

从这里开始，您可以像使用 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，它将把这些输入转发到我们的实现，以记录当输入流经它时发生的操作序列。这样做的结果实际上是 eager 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 代码中就这么简单！

将自定义算子与 Script 一起使用

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

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

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

要将此函数从原始 PyTorch 转换为 TorchScript，我们使用 @torch.jit.script 对其进行注解

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

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

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

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

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

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

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

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

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

注意

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

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

在 C++ 中使用 TorchScript 自定义算子

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

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

注意

从技术上讲，您也可以在运行时将共享库动态加载到您的 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().
  torch::jit::script::Module module = torch::jit::load(argv[1]);

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

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

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

以及一个小的 CMakeLists.txt 文件

cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
project(example_app)

find_package(Torch REQUIRED)

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

此时，我们应该能够构建应用程序

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

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

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

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

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

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

compute.save("example.pt")

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

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

或者可能不行。 也许还不行。 当然！ 我们还没有将自定义算子库与我们的应用程序链接。 让我们现在就这样做，为了正确地做到这一点，让我们稍微更新一下我们的文件组织，使其看起来像这样

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

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

cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
project(example_app)

find_package(Torch REQUIRED)

add_subdirectory(warp_perspective)

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

这个基本的 CMake 配置看起来与之前非常相似，除了我们将 warp_perspective CMake 构建添加为子目录。 一旦其 CMake 代码运行，我们将我们的 example_app 应用程序与 warp_perspective 共享库链接。

注意

上述示例中嵌入了一个关键细节：-Wl,--no-as-needed 前缀到 warp_perspective 链接行。 这是必需的，因为我们实际上不会在我们的应用程序代码中调用 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++ 之间界限平滑融合的用例。

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

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

“构建自定义算子”部分解释了如何使用 CMake 将自定义算子构建到共享库中。 本附录概述了另外两种编译方法。 它们都使用 Python 作为编译过程的“驱动程序”或“接口”。 此外，两者都重用了 PyTorch 为 *C++ 扩展* 提供的 现有基础设施，C++ 扩展是 TorchScript 自定义算子的原始（eager）PyTorch 等效项，它依赖于 pybind11 来实现从 C++ 到 Python 的函数“显式”绑定。

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

使用 JIT 编译构建

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

注意

此处的“JIT 编译”与 TorchScript 编译器中发生的 JIT 编译无关，后者用于优化您的程序。 它只是意味着您的自定义算子 C++ 代码将在您第一次导入它时在系统 /tmp 目录下的文件夹中编译，就像您事先自己编译过它一样。

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

import torch.utils.cpp_extension

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

print(torch.ops.my_ops.warp_perspective)

这应该大致打印

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

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

import torch
import torch.utils.cpp_extension

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

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

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

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

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

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

print(torch.ops.my_ops.warp_perspective)

自然地，最佳实践是仅当您的源代码相当短时才使用 torch.utils.cpp_extension.load_inline。

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

使用 Setuptools 构建

从 Python 专门构建我们的自定义算子的第二种方法是使用 setuptools。 这具有以下优点：setuptools 具有非常强大和广泛的接口，用于构建用 C++ 编写的 Python 模块。 但是，由于 setuptools 实际上旨在构建 Python 模块，而不是普通的共享库（这些共享库不具有 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>