扩展 TorchScript 以使用自定义 C++运算符
创建于：2025 年 4 月 1 日 | 最后更新：2025 年 4 月 1 日 | 最后验证：2024 年 11 月 5 日
警告
本教程自 PyTorch 2.4 版本起已过时。请参阅 PyTorch 自定义运算符以获取最新、最全面的指南。
PyTorch 1.0 版本引入了 PyTorch 的新编程模型，称为 TorchScript。TorchScript 是 Python 编程语言的一个子集，可以被 TorchScript 编译器解析、编译和优化。此外，编译后的 TorchScript 模型可以选择序列化为磁盘文件格式，您可以从纯 C++（以及 Python）中加载并运行这些文件进行推理。
TorchScript 支持 torch 包提供的操作的大子集，允许您将许多类型的复杂模型纯粹表达为 PyTorch“标准库”中的一系列张量操作。尽管如此，有时您可能需要扩展 TorchScript 以使用自定义的 C++或 CUDA 函数。虽然我们建议您只有在您的想法无法（足够高效地）作为简单的 Python 函数表达时才使用此选项，但我们确实提供了一个非常友好且简单的接口，用于使用 ATen（PyTorch 的高性能 C++张量库）定义自定义的 C++和 CUDA 内核。一旦绑定到 TorchScript，您可以将这些自定义内核（或“操作”）嵌入到您的 TorchScript 模型中，并在 Python 中以及直接在 C++中以序列化形式执行它们。
以下段落提供了一个编写 TorchScript 自定义操作符以调用 OpenCV（一个用 C++编写的计算机视觉库）的示例。我们将讨论如何在 C++中处理张量，如何高效地将它们转换为第三方张量格式（在本例中为 OpenCV Mat ），如何将您的操作符注册到 TorchScript 运行时，以及最后如何在 Python 和 C++中编译操作符并使用它。
在 C++中实现自定义操作符
在本教程中，我们将公开 OpenCV 的 warpPerspective 函数，该函数将图像应用透视变换，并将其作为自定义操作符暴露给 TorchScript。第一步是编写我们的自定义操作符的实现。让我们将这个实现的文件命名为 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 中进行研究和实验，或者在没有 Python 的环境中将其构建到 C++中进行推理。构建我们的操作符有多种方法，可以使用纯 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.<namespace>.<function> 引用已加载的操作符，其中 <namespace> 是操作符名称的命名空间部分， <function> 是操作符的函数名。对于我们上面编写的操作符，命名空间是 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 trace 中，就像它是 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 自定义操作集成到 traced 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 的引用时，它将找到我们通过 TORCH_LIBRARY 函数在 C++中注册的实现，并将其编译成其图表示：
>>> 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 自定义算子
火炬脚本的一个有用特性是将模型序列化到磁盘文件中。这个文件可以通过网络发送，存储在文件系统中，或者更重要的是，可以动态反序列化并执行，而无需保留原始源代码。这在 Python 中是可能的，同样在 C++中也是。为此，PyTorch 提供了一个用于反序列化和执行火炬脚本模型的纯 C++ API。如果您还没有阅读，请阅读有关在 C++中加载和运行序列化火炬脚本模型的教程，接下来的几段将基于此进行。
简而言之，自定义操作符可以像常规操作符一样执行，即使是从文件反序列化并在 C++中运行。为此，唯一的要求是将我们之前构建的自定义操作符共享库与执行模型的 C++应用程序链接起来。在 Python 中，这可以通过调用操作符简单地完成。在 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 共享库链接起来。
注意
上面的例子中嵌入了一个关键细节： -Wl,--no-as-needed 前缀到 warp_perspective 链接行的 -Wl,--no-as-needed。这是必需的，因为我们实际上不会在我们的应用程序代码中调用 warp_perspective 共享库中的任何函数。我们只需要运行 TORCH_LIBRARY 函数。不方便的是，这会让链接器误以为它可以完全跳过对库的链接。在 Linux 上， -Wl,--no-as-needed 标志强制进行链接（注意：此标志是 Linux 特有的！）。还有其他解决方案。最简单的是在操作员库中定义一些需要从主应用程序中调用的函数。这可以像在某个头文件中声明的 void init(); 函数一样简单，然后在操作员库中定义为 void init() { } 。在主应用程序中调用这个 init() 函数会给链接器一个印象，即这是一个值得链接的库。不幸的是，这超出了我们的控制范围，我们宁愿让您知道原因和简单的解决方案，而不是给您一些难以理解的宏来插入到您的代码中。
现在，由于我们发现在顶级目录中找到了 Torch 包， CMakeLists.txt 文件在 warp_perspective 子目录中的长度可以缩短一些。它应该看起来像这样：
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 自定义操作符的 vanilla（急切）PyTorch 等效，这些操作符依赖于 pybind11 将 C++中的函数“显式”绑定到 Python。
第一种方法使用 C++扩展方便的即时（JIT）编译接口，在您第一次运行 PyTorch 脚本时在后台编译您的代码。第二种方法依赖于古老的 setuptools 包，并涉及编写一个单独的 setup.py 文件。这允许更高级的配置以及与其他 setuptools -based 项目的集成。以下我们将详细介绍这两种方法。
使用即时编译构建
PyTorch C++扩展工具包提供的即时编译功能允许将自定义算子的编译直接嵌入到 Python 代码中，例如在训练脚本的开头。
备注
这里的“即时编译”与在 TorchScript 编译器中进行的即时编译优化程序无关。它仅仅意味着，当您第一次导入它时，您的自定义算子 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 一样：
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>
第二种即时编译形式允许你将自定义 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 从模块期望的必要入口点），这条路径可能会有点奇怪。话虽如此，您只需要在 CMakeLists.txt 的位置放置一个 setup.py 文件，它看起来像这样：
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>
