使用自定义 C++ 操作扩展 TorchScript
本教程自 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++ 中实现自定义操作符
在本教程中，我们将把 OpenCV 中的 warpPerspective 函数（该函数对图像应用透视变换）作为自定义操作符暴露给 TorchScript。第一步是使用 C++ 编写自定义操作符的实现。我们将这个实现文件命名为 op.cpp，并使其内容如下：
torch::Tensorwarp_perspective(torch::Tensorimage,torch::Tensorwarp){
// BEGIN image_mat
cv::Matimage_mat(/*rows=*/image.size(0),
/*cols=*/image.size(1),
/*type=*/CV_32FC1,
/*data=*/image.data_ptr<float>());
// END image_mat
// BEGIN warp_mat
cv::Matwarp_mat(/*rows=*/warp.size(0),
/*cols=*/warp.size(1),
/*type=*/CV_32FC1,
/*data=*/warp.data_ptr<float>());
// END warp_mat
// BEGIN output_mat
cv::Matoutput_mat;
cv::warpPerspective(image_mat,output_mat,warp_mat,/*dsize=*/{8,8});
// END output_mat
// BEGIN output_tensor
torch::Tensoroutput=torch::from_blob(output_mat.ptr<float>(),/*sizes=*/{8,8});
returnoutput.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::Matimage_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::Matwarp_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::Matoutput_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::Tensoroutput=torch::from_blob(output_mat.ptr<float>(),/*sizes=*/{8,8});
returnoutput.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(VERSION3.1FATAL_ERROR)
project(warp_perspective)
find_package(TorchREQUIRED)
find_package(OpenCVREQUIRED)
# Define our library target
add_library(warp_perspectiveSHAREDop.cpp)
# Enable C++14
target_compile_features(warp_perspectivePRIVATEcxx_std_14)
# Link against LibTorch
target_link_libraries(warp_perspective"${TORCH_LIBRARIES}")
# Link against OpenCV
target_link_libraries(warp_perspectiveopencv_coreopencv_imgproc)
要构建我们的操作符，可以从 warp_perspective 文件夹中运行以下命令：
$mkdirbuild
$cdbuild
$cmake-DCMAKE_PREFIX_PATH="$(python-c'import torch.utils; print(torch.utils.cmake_prefix_path)')"..
*-TheCcompileridentificationisGNU5.4.0
*-TheCXXcompileridentificationisGNU5.4.0
*-CheckforworkingCcompiler:/usr/bin/cc
*-CheckforworkingCcompiler:/usr/bin/cc--works
*-DetectingCcompilerABIinfo
*-DetectingCcompilerABIinfo-done
*-DetectingCcompilefeatures
*-DetectingCcompilefeatures-done
*-CheckforworkingCXXcompiler:/usr/bin/c++
*-CheckforworkingCXXcompiler:/usr/bin/c++--works
*-DetectingCXXcompilerABIinfo
*-DetectingCXXcompilerABIinfo-done
*-DetectingCXXcompilefeatures
*-DetectingCXXcompilefeatures-done
*-Lookingforpthread.h
*-Lookingforpthread.h-found
*-Lookingforpthread_create
*-Lookingforpthread_create-notfound
*-Lookingforpthread_createinpthreads
*-Lookingforpthread_createinpthreads-notfound
*-Lookingforpthread_createinpthread
*-Lookingforpthread_createinpthread-found
*-FoundThreads:TRUE
*-Foundtorch:/libtorch/lib/libtorch.so
*-Configuringdone
*-Generatingdone
*-Buildfileshavebeenwrittento:/warp_perspective/build
$make-j
Scanningdependenciesoftargetwarp_perspective
[50%]BuildingCXXobjectCMakeFiles/warp_perspective.dir/op.cpp.o
[100%]LinkingCXXsharedlibrarylibwarp_perspective.so
[100%]Builttargetwarp_perspective
这将会在 build 文件夹中生成一个 libwarp_perspective.so 共享库文件。在上述 cmake 命令中，我们使用了辅助变量 torch.utils.cmake_prefix_path，以便方便地告知我们 PyTorch 安装的 cmake 文件所在的位置。
我们将在下面详细探讨如何使用和调用我们的算子，但为了提前感受成功的喜悦，可以尝试在 Python 中运行以下代码：
importtorch
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 模块中使用，但我们也可以在普通的 eager 模式下使用 PyTorch，并传递常规的 PyTorch 张量给它：
importtorch
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 代码开始：
defcompute(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 trace 中，就像使用 torch.relu 或其他任何 torch 函数一样：
defcompute(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 而非追踪（tracing）的一个特别原因是，追踪无法捕获 PyTorch 代码中的控制流。因此，让我们考虑一下这个使用了控制流的函数：
defcompute(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
defcompute(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
defcompute(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 模型的教程，接下来的几段内容将在此基础上展开。
简而言之，自定义运算符可以像常规的 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>
intmain(intargc,constchar*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::Modulemodule=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::Tensoroutput=module.forward(std::move(inputs)).toTensor();
std::cout<<output<<std::endl;
}
伴随一个小的 CMakeLists.txt 文件：
cmake_minimum_required(VERSION3.1FATAL_ERROR)
project(example_app)
find_package(TorchREQUIRED)
add_executable(example_appmain.cpp)
target_link_libraries(example_app"${TORCH_LIBRARIES}")
target_compile_features(example_appPRIVATEcxx_range_for)
此时，我们应该能够构建应用程序：
$mkdirbuild
$cdbuild
$cmake-DCMAKE_PREFIX_PATH="$(python-c'import torch.utils; print(torch.utils.cmake_prefix_path)')"..
*-TheCcompileridentificationisGNU5.4.0
*-TheCXXcompileridentificationisGNU5.4.0
*-CheckforworkingCcompiler:/usr/bin/cc
*-CheckforworkingCcompiler:/usr/bin/cc--works
*-DetectingCcompilerABIinfo
*-DetectingCcompilerABIinfo-done
*-DetectingCcompilefeatures
*-DetectingCcompilefeatures-done
*-CheckforworkingCXXcompiler:/usr/bin/c++
*-CheckforworkingCXXcompiler:/usr/bin/c++--works
*-DetectingCXXcompilerABIinfo
*-DetectingCXXcompilerABIinfo-done
*-DetectingCXXcompilefeatures
*-DetectingCXXcompilefeatures-done
*-Lookingforpthread.h
*-Lookingforpthread.h-found
*-Lookingforpthread_create
*-Lookingforpthread_create-notfound
*-Lookingforpthread_createinpthreads
*-Lookingforpthread_createinpthreads-notfound
*-Lookingforpthread_createinpthread
*-Lookingforpthread_createinpthread-found
*-FoundThreads:TRUE
*-Foundtorch:/libtorch/lib/libtorch.so
*-Configuringdone
*-Generatingdone
*-Buildfileshavebeenwrittento:/example_app/build
$make-j
Scanningdependenciesoftargetexample_app
[50%]BuildingCXXobjectCMakeFiles/example_app.dir/main.cpp.o
[100%]LinkingCXXexecutableexample_app
[100%]Builttargetexample_app
然后先不传递模型直接运行它：
$./example_app
usage:example_app<path-to-exported-script-module>
接下来，让我们序列化之前编写的脚本函数，该函数使用了我们的自定义操作符：
torch.ops.load_library("libwarp_perspective.so")
@torch.jit.script
defcompute(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_appexample.pt
terminatecalledafterthrowinganinstanceof'torch::jit::script::ErrorReport'
what():
Schemanotfoundfornode.Fileabugreport.
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(VERSION3.1FATAL_ERROR)
project(example_app)
find_package(TorchREQUIRED)
add_subdirectory(warp_perspective)
add_executable(example_appmain.cpp)
target_link_libraries(example_app"${TORCH_LIBRARIES}")
target_link_libraries(example_app-Wl,--no-as-neededwarp_perspective)
target_compile_features(example_appPRIVATEcxx_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(OpenCVREQUIRED)
add_library(warp_perspectiveSHAREDop.cpp)
target_compile_features(warp_perspectivePRIVATEcxx_range_for)
target_link_libraries(warp_perspectivePRIVATE"${TORCH_LIBRARIES}")
target_link_libraries(warp_perspectivePRIVATEopencv_coreopencv_photo)
让我们重新构建示例应用程序，它还将与自定义操作符库链接。在顶层的 example_app 目录中：
$mkdirbuild
$cdbuild
$cmake-DCMAKE_PREFIX_PATH="$(python-c'import torch.utils; print(torch.utils.cmake_prefix_path)')"..
*-TheCcompileridentificationisGNU5.4.0
*-TheCXXcompileridentificationisGNU5.4.0
*-CheckforworkingCcompiler:/usr/bin/cc
*-CheckforworkingCcompiler:/usr/bin/cc--works
*-DetectingCcompilerABIinfo
*-DetectingCcompilerABIinfo-done
*-DetectingCcompilefeatures
*-DetectingCcompilefeatures-done
*-CheckforworkingCXXcompiler:/usr/bin/c++
*-CheckforworkingCXXcompiler:/usr/bin/c++--works
*-DetectingCXXcompilerABIinfo
*-DetectingCXXcompilerABIinfo-done
*-DetectingCXXcompilefeatures
*-DetectingCXXcompilefeatures-done
*-Lookingforpthread.h
*-Lookingforpthread.h-found
*-Lookingforpthread_create
*-Lookingforpthread_create-notfound
*-Lookingforpthread_createinpthreads
*-Lookingforpthread_createinpthreads-notfound
*-Lookingforpthread_createinpthread
*-Lookingforpthread_createinpthread-found
*-FoundThreads:TRUE
*-Foundtorch:/libtorch/lib/libtorch.so
*-Configuringdone
*-Generatingdone
*-Buildfileshavebeenwrittento:/warp_perspective/example_app/build
$make-j
Scanningdependenciesoftargetwarp_perspective
[25%]BuildingCXXobjectwarp_perspective/CMakeFiles/warp_perspective.dir/op.cpp.o
[50%]LinkingCXXsharedlibrarylibwarp_perspective.so
[50%]Builttargetwarp_perspective
Scanningdependenciesoftargetexample_app
[75%]BuildingCXXobjectCMakeFiles/example_app.dir/main.cpp.o
[100%]LinkingCXXexecutableexample_app
[100%]Builttargetexample_app
如果我们现在运行 example_app 可执行文件并传入我们序列化的模型，应该会得到一个圆满的结果：
$./example_appexample.pt
11.41255.82629.53458.611112.3997
7.468313.59699.085011.06989.4008
7.459715.092612.57278.93199.0666
9.483411.17479.016210.95218.6269
10.000010.000010.000010.000010.0000
10.000010.000010.000010.000010.0000
10.000010.000010.000010.000010.0000
10.000010.000010.000010.000010.0000
[Variable[CPUFloatType]{8,5}]
成功！您现在可以开始推理了。
总结
本教程向您详细介绍了如何用 C++ 实现自定义的 TorchScript 运算符，如何将其构建为共享库，如何在 Python 中使用它来定义 TorchScript 模型，以及最后如何将其加载到 C++ 应用程序中以进行推理工作。现在，您已经准备好通过 C++ 运算符来扩展您的 TorchScript 模型，这些运算符可以与第三方 C++ 库进行交互，编写自定义的高性能 CUDA 内核，或实现任何需要 Python、TorchScript 和 C++ 之间无缝集成的用例。
如往常一样，如果您遇到任何问题或有疑问，可以使用我们的论坛或GitHub issues与我们联系。此外，我们的常见问题解答 (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 目录下的一个文件夹中进行编译，就像您事先自己编译了一样。
这个即时编译功能有两种形式。在第一种形式中，您仍然将操作符的实现放在一个单独的文件中（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 所做的那样：
importtorch.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：
importtorch
importtorch.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，它看起来像这样：
fromsetuptoolsimport setup
fromtorch.utils.cpp_extensionimport 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。
如果现在我们在终端中运行 python setup.py build develop，并且 setup.py 所在的文件夹中，我们应该会看到类似以下内容：
$pythonsetup.pybuilddevelop
runningbuild
runningbuild_ext
building'warp_perspective'extension
creatingbuild
creatingbuild/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-cop.cpp-obuild/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:commandlineoption‘-Wstrict-prototypes’isvalidforC/ObjCbutnotforC++
creatingbuild/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-obuild/lib.linux-x86_64-3.7/warp_perspective.so
runningdevelop
runningegg_info
creatingwarp_perspective.egg-info
writingwarp_perspective.egg-info/PKG-INFO
writingdependency_linkstowarp_perspective.egg-info/dependency_links.txt
writingtop-levelnamestowarp_perspective.egg-info/top_level.txt
writingmanifestfile'warp_perspective.egg-info/SOURCES.txt'
readingmanifestfile'warp_perspective.egg-info/SOURCES.txt'
writingmanifestfile'warp_perspective.egg-info/SOURCES.txt'
runningbuild_ext
copyingbuild/lib.linux-x86_64-3.7/warp_perspective.so->
Creating/root/local/miniconda/lib/python3.7/site-packages/warp-perspective.egg-link(linkto.)
Addingwarp-perspective0.0.0toeasy-install.pthfile
Installed/warp_perspective
Processingdependenciesforwarp-perspective==0.0.0
Finishedprocessingdependenciesforwarp-perspective==0.0.0
这将生成一个名为 warp_perspective.so 的共享库，我们可以像之前一样将其传递给 torch.ops.load_library，以使我们的操作符对 TorchScript 可见：
>>> importtorch
>>> 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>