• Tutorials >
  • Extending TorchScript with Custom C++ Operators
Shortcuts

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

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

警告

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

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

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

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

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

在本教程中,我们将从OpenCV向TorchScript公开warpPerspective函数,该函数对图像应用透视变换,作为自定义操作符。第一步是用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库的更多信息,请参阅此说明,该库为PyTorch提供了Tensor类。此外,本教程描述了如何在C++中分配和初始化新的张量对象(此操作符不需要)。

注意

TorchScript 编译器理解固定数量的类型。只有这些类型可以作为自定义运算符的参数。目前这些类型是:torch::Tensor, torch::Scalar, double, int64_t 和这些类型的 std::vector。请注意,支持 double 而不支持 float,并且 支持 int64_t 而不支持其他整数类型,如 int, shortlong

在我们的函数内部,我们需要做的第一件事是将我们的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。为此,我们将OpenCV函数传递给image_matwarp_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();

我们在OpenCV的Mat类上使用.ptr()方法来获取底层数据的原始指针(就像之前PyTorch张量的.data_ptr()一样)。我们还指定了张量的输出形状,我们将其硬编码为8 x 8torch::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,也可以使用像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

这将在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.., 其中 是你操作符名称的命名空间部分,而 是你操作符的函数名称。对于我们上面编写的操作符, 命名空间是 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 模型的教程,接下来的几段将在此基础上展开。

简而言之,自定义操作符可以像常规的torch操作符一样执行,即使是从文件中反序列化并在C++中运行。唯一的要求是将我们之前构建的自定义操作符共享库与执行模型的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++操作符扩展您的TorchScript模型,这些操作符与第三方C++库接口,编写自定义高性能CUDA内核,或实现任何其他需要Python、TorchScript和C++之间无缝融合的用例。

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

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

“构建自定义操作符”部分解释了如何使用CMake将自定义操作符构建到共享库中。本附录概述了两种进一步的编译方法。这两种方法都使用Python作为编译过程的“驱动程序”或“接口”。此外,两者都重用了PyTorch为*C++扩展*提供的现有基础设施,这些扩展是TorchScript自定义操作符的普通(急切)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中使用此功能,不应多次执行包含注册的单元格,因为每次执行都会注册一个新库并重新注册自定义操作符。如果需要重新执行,请事先重启您的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>