使用自定义C++操作符扩展TorchScript¶
Created On: Nov 28, 2018 | Last Updated: Jul 22, 2024 | Last Verified: Nov 05, 2024
警告
本教程自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 （一个用C++编写的计算机视觉库）的示例。我们将讨论如何在C++中处理张量，如何高效地将它们转换为第三方张量格式（在本例中为OpenCV Mat），如何将操作符注册到TorchScript运行时，以及如何编译操作符并在Python和C++中使用它。
在C++中实现自定义操作符¶
在本教程中，我们将会将 warpPerspective 函数从OpenCV暴露为TorchScript的自定义操作符，该函数对图像应用透视变换。第一步是使用C++编写我们自定义操作符的实现。我们将这个实现的文件命名为 op.cpp，并让其代码如下所示：
torch::Tensor warp_perspective(torch::Tensor image, torch::Tensor warp) {
// BEGIN image_mat
cv::Mat image_mat(/*rows=*/image.size(0),
/*cols=*/image.size(1),
/*type=*/CV_32FC1,
/*data=*/image.data_ptr<float>());
// END image_mat
// BEGIN warp_mat
cv::Mat warp_mat(/*rows=*/warp.size(0),
/*cols=*/warp.size(1),
/*type=*/CV_32FC1,
/*data=*/warp.data_ptr<float>());
// END warp_mat
// BEGIN output_mat
cv::Mat output_mat;
cv::warpPerspective(image_mat, output_mat, warp_mat, /*dsize=*/{8, 8});
// END output_mat
// BEGIN output_tensor
torch::Tensor output = torch::from_blob(output_mat.ptr<float>(), /*sizes=*/{8, 8});
return output.clone();
// END output_tensor
}
该操作符的代码非常简短。在文件的顶部，我们引入了OpenCV的头文件``opencv2/opencv.hpp``，以及``torch/script.h``头文件，该文件包含了必要的PyTorch C++ API功能，可以用来编写自定义TorchScript操作符。我们的函数``warp_perspective``接受两个参数：输入的``image``和我们希望应用到图像上的``warp``变换矩阵。这些输入的类型是``torch::Tensor``，即PyTorch在C++中的张量类型（也是Python中所有张量的底层类型）。我们的``warp_perspective``函数的返回类型也是``torch::Tensor``。
注意
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。为此，我们向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();
我们在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``函数实际上做了很多工作：它使用模板元编程来检查您函数的类型签名，并将其翻译为一个操作符模式（schema），该模式指定了操作符在TorchScript类型系统中的类型。
构建自定义操作符¶
现在我们已经在C++中实现了自定义操作符并编写了其注册代码，是时候将操作符构建为（共享）库了，这样我们就可以在Python中加载它用于研究和实验，或者在无Python环境中用C++进行推理了。存在多种构建操作符的方法，可以使用纯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 自定义运算符是由 PyTorch 动态绑定的。pybind11 在绑定类型和类到 Python 时提供更多灵活性，因此更适合纯粹的 eager 代码，但不支持 TorchScript 运算符。
从这里开始，你可以像使用 torch 包中的其他函数一样，在脚本化或跟踪代码中使用自定义运算符。事实上，像 torch.matmul 这样的“标准库”函数基本上通过与自定义运算符相同的注册路径，因此在 TorchScript 中使用自定义运算符就像使用其他函数一样（不过，标准库函数拥有自定义编写的 Python 参数解析逻辑，这与 torch.ops 的参数解析有所不同）。
在跟踪中使用自定义运算符¶
让我们从将自定义运算符嵌入到跟踪函数开始。回顾一下，对于跟踪，我们从一些普通的 PyTorch 代码开始：
def compute(x, y, z):
return x.matmul(y) + torch.relu(z)
然后调用 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)
现在令人兴奋的发现是，我们可以像使用 torch.relu 或任何其他 torch 函数一样，轻松将自定义运算符直接插入 PyTorch 跟踪中：
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 导入包含运算符的库，并像在跟踪或脚本化的 TorchScript 代码中调用其他 torch 运算符一样调用自定义运算符。
在 C++ 中使用 TorchScript 自定义运算符¶
TorchScript 的一个有用功能是可以将模型序列化为磁盘文件。这个文件可以通过网络发送、存储在文件系统中，或者更重要的是可以动态反序列化并执行，而无需保留原始源代码。这在 Python 中适用，也在 C++ 中适用。为此，PyTorch 提供了 一个纯 C++ API 用于反序列化和执行 TorchScript 模型。如果你还没有，请阅读 加载和运行序列化 TorchScript 模型的 C++ 教程，接下来的几段解释的是基于该内容的。
简而言之，即使从文件中反序列化并在 C++ 中运行，自定义运算符也可以像常规 torch 运算符一样执行。这唯一的要求是将之前构建的自定义运算符共享库与执行模型的 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 目标作为我们应用程序目标的一个子目录。顶级 example_app 文件夹中的 CMakeLists.txt 文件应如下所示：
cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
project(example_app)
find_package(Torch REQUIRED)
add_subdirectory(warp_perspective)
add_executable(example_app main.cpp)
target_link_libraries(example_app "${TORCH_LIBRARIES}")
target_link_libraries(example_app -Wl,--no-as-needed warp_perspective)
target_compile_features(example_app PRIVATE cxx_range_for)
这个基本的 CMake 配置看起来与之前很相似，但我们添加了 warp_perspective CMake 构建作为一个子目录。一旦它的 CMake 代码运行完，我们就将 example_app 应用程序与 warp_perspective 共享库链接起来。
注意
上面示例中嵌入了一个重要细节：warp_perspective 链接行的 -Wl,--no-as-needed 前缀。这是必要的，因为我们的应用程序代码中实际上不会调用 warp_perspective 共享库中的任何函数。我们只需要运行 TORCH_LIBRARY 函数。不幸的是，这会迷惑链接器，使其认为可以完全跳过链接到这个库。在 Linux 上，-Wl,--no-as-needed 标志强制进行链接（注意：此标志仅适用于 Linux！）。还有其他一些解决方法。最简单的是在操作符库中定义 某些函数，以便需要在主应用程序中调用。这可以是一个简单的声明在某些头文件中的函数 void init(); ，然后在运算符库中定义为 void init() { }。在主应用程序中调用这个 init() 函数会让链接器认为这是一个值得链接的库。不幸的是，我们对此无能为力，所以我们更愿意告诉你背后的原因和这个简单的解决办法，而不是给你一些黑箱宏放到代码中。
现在，由于我们在顶级找到了 Torch 包，因此在 warp_perspective 子目录中的 CMakeLists.txt 文件可以稍加简化，如下所示：
find_package(OpenCV REQUIRED)
add_library(warp_perspective SHARED op.cpp)
target_compile_features(warp_perspective PRIVATE cxx_range_for)
target_link_libraries(warp_perspective PRIVATE "${TORCH_LIBRARIES}")
target_link_libraries(warp_perspective PRIVATE opencv_core opencv_photo)
让我们重新构建示例应用程序，它还会链接自定义运算符库。在顶级 example_app 目录中：
$ mkdir build
$ cd build
$ cmake -DCMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" ..
-- The C compiler identification is GNU 5.4.0
-- The CXX compiler identification is GNU 5.4.0
-- Check for working C compiler: /usr/bin/cc
-- Check for working C compiler: /usr/bin/cc -- works
-- Detecting C compiler ABI info
-- Detecting C compiler ABI info - done
-- Detecting C compile features
-- Detecting C compile features - done
-- Check for working CXX compiler: /usr/bin/c++
-- Check for working CXX compiler: /usr/bin/c++ -- works
-- Detecting CXX compiler ABI info
-- Detecting CXX compiler ABI info - done
-- Detecting CXX compile features
-- Detecting CXX compile features - done
-- Looking for pthread.h
-- Looking for pthread.h - found
-- Looking for pthread_create
-- Looking for pthread_create - not found
-- Looking for pthread_create in pthreads
-- Looking for pthread_create in pthreads - not found
-- Looking for pthread_create in pthread
-- Looking for pthread_create in pthread - found
-- Found Threads: TRUE
-- Found torch: /libtorch/lib/libtorch.so
-- Configuring done
-- Generating done
-- Build files have been written to: /warp_perspective/example_app/build
$ make -j
Scanning dependencies of target warp_perspective
[ 25%] Building CXX object warp_perspective/CMakeFiles/warp_perspective.dir/op.cpp.o
[ 50%] Linking CXX shared library libwarp_perspective.so
[ 50%] Built target warp_perspective
Scanning dependencies of target example_app
[ 75%] Building CXX object CMakeFiles/example_app.dir/main.cpp.o
[100%] Linking CXX executable example_app
[100%] Built target example_app
如果我们现在运行 example_app 二进制文件并将其与序列化模型连接，我们应该会得到一个满意的结果：
$ ./example_app example.pt
11.4125 5.8262 9.5345 8.6111 12.3997
7.4683 13.5969 9.0850 11.0698 9.4008
7.4597 15.0926 12.5727 8.9319 9.0666
9.4834 11.1747 9.0162 10.9521 8.6269
10.0000 10.0000 10.0000 10.0000 10.0000
10.0000 10.0000 10.0000 10.0000 10.0000
10.0000 10.0000 10.0000 10.0000 10.0000
10.0000 10.0000 10.0000 10.0000 10.0000
[ Variable[CPUFloatType]{8,5} ]
成功！现在你可以开始推理。
总结¶
本教程向你展示了如何在 C++ 中实现一个自定义 TorchScript 运算符，如何将其构建为共享库，如何在 Python 中使用它来定义 TorchScript 模型以及最后如何将其加载到一个 C++ 应用程序中以进行推理工作负载。现在，你已经准备好扩展你的 TorchScript 模型，使用能够与第三方 C++ 库交互的 C++ 运算符，编写自定义高性能 CUDA 内核，或实现任何需要在 Python、TorchScript 和 C++ 之间平滑融合的用例。
一如既往，如果您遇到任何问题或有任何疑问，可以使用我们的 论坛 或 GitHub问题 联系我们。此外，我们的 常见问题解答 (FAQ) 页 可能也有有用的信息。
附录 A: 构建自定义算子的更多方法¶
章节“构建自定义算子”讲解了如何使用 CMake 将自定义算子构建为一个共享库。本附录概述了另外两种编译方法。这两种方法都使用 Python 作为编译过程的“驱动程序”或“接口”。此外，两种方法都重新利用了 PyTorch 提供的现有基础设施，参考`C++ 扩展 <https://pytorch.org/docs/stable/cpp_extension.html>`_，这实际是 TorchScript 自定义算子在原生（动态）PyTorch 中的等价形式，使用 pybind11 来将 C++ 的功能“显式地”绑定到 Python。
第一种方法使用 C++ 扩展工具的`即时 (JIT) 编译接口 <https://pytorch.org/docs/stable/cpp_extension.html#torch.utils.cpp_extension.load>`_ 来在首次运行时于 PyTorch 脚本后台编译你的代码。第二种方法依赖传统的 setuptools 包，需要编写一个独立的 setup.py 文件。这种方法允许更高级的配置，并可与基于 setuptools 的其它项目集成。我们将在下面详细探索这两种方法。
使用 JIT 编译构建¶
PyTorch C++ 扩展工具提供的 JIT 编译功能允许直接在 Python 代码（比如你的训练脚本开头）中嵌入自定义算子的编译过程。
备注
这里的 “JIT 编译” 并非指 TorchScript 编译器中的 JIT 编译用于优化程序，而是表示自定义算子的 C++ 代码将在首次导入时于系统的 /tmp 目录下一个文件夹中被编译，就像你事先自己编译它一样。
此 JIT 编译功能有两种形式。第一种形式，你将算子实现保留在单独的文件中（例如 op.cpp），然后使用 torch.utils.cpp_extension.load() 编译扩展。通常，这个函数返回一个 Python 模块，暴露你的 C++ 扩展。然而，由于我们不打算将自定义算子编译成一个独立的 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>
