使用自定义 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 自定义算子来调用用 C++ 编写的计算机视觉库 OpenCV。我们将讨论如何在 C++ 中处理张量，如何有效地将它们转换为第三方张量格式（在本例中是 OpenCV Mat），如何向 TorchScript 运行时注册你的算子，最后是如何编译算子并在 Python 和 C++ 中使用它。
在 C++ 中实现自定义算子¶
在本教程中，我们将 OpenCV 的 warpPerspective 函数（该函数对图像应用透视变换）作为自定义算子暴露给 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 构造函数的特殊之处在于它不会复制输入数据。相反，它只会在所有对 image_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 库，我们的注册语法与其非常相似。要注册单个函数，我们在 op.cpp 文件的顶层某处写入：
TORCH_LIBRARY(my_ops, m) {
m.def("warp_perspective", warp_perspective);
}
TORCH_LIBRARY 宏创建一个函数，该函数会在程序启动时被调用。你的库的名称（my_ops）作为第一个参数给出（不应加引号）。第二个参数（m）定义了一个 torch::Library 类型的变量，它是注册算子的主要接口。方法 Library::def 实际上创建了一个名为 warp_perspective 的算子，并将其暴露给 Python 和 TorchScript。你可以通过多次调用 def 来定义任意数量的算子。
在幕后，def 函数实际上做了不少工作：它使用模板元编程来检查函数的类型签名，并将其转换为一个算子 schema，该 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 使用。虽然此函数可以在 scripted 或 traced 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 算子。
从现在开始，你可以在 scripted 或 traced 代码中使用你的自定义算子，就像使用 torch 包中的其他函数一样。事实上，像 torch.matmul 这样的“标准库”函数很大程度上遵循与自定义算子相同的注册路径，这使得自定义算子在 TorchScript 中的使用方式和位置方面真正成为一等公民。（然而，一个区别是标准库函数有自定义编写的 Python 参数解析逻辑，这与 torch.ops 参数解析不同。）
使用 Tracing 自定义算子¶
让我们从将我们的算子嵌入 traced 函数开始。回想一下，对于 tracing，我们从一些普通 PyTorch 代码开始：
def compute(x, y, z):
return x.matmul(y) + torch.relu(z)
然后在其上调用 torch.jit.trace。我们还会向 torch.jit.trace 传递一些示例输入，它将把这些输入转发给我们的实现，以记录输入流经时发生的操作序列。其结果实际上是一个 eager PyTorch 程序的“冻结”版本，TorchScript 编译器可以对其进行进一步分析、优化和序列化：
inputs = [torch.randn(4, 8), torch.randn(8, 5), torch.randn(4, 5)]
trace = torch.jit.trace(compute, inputs)
print(trace.graph)
生成：
graph(%x : Float(4:8, 8:1),
%y : Float(8:5, 5:1),
%z : Float(4:5, 5:1)):
%3 : Float(4:5, 5:1) = aten::matmul(%x, %y) # test.py:10:0
%4 : Float(4:5, 5:1) = aten::relu(%z) # test.py:10:0
%5 : int = prim::Constant[value=1]() # test.py:10:0
%6 : Float(4:5, 5:1) = aten::add(%3, %4, %5) # test.py:10:0
return (%6)
现在，令人兴奋的发现是，我们可以简单地将自定义算子放入 PyTorch 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)
然后像之前一样进行 tracing：
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 代码中就是如此简单！
使用 Script 自定义算子¶
除了 tracing 之外，获得 PyTorch 程序 TorchScript 表示的另一种方法是直接**在** TorchScript 中编写代码。TorchScript 很大程度上是 Python 语言的一个子集，带有一些限制，这些限制使 TorchScript 编译器更容易理解程序。你通过为自由函数添加 @torch.jit.script 注解，以及为类中的方法添加 @torch.jit.script_method 注解（这些类也必须派生自 torch.jit.ScriptModule）来将你的普通 PyTorch 代码转换为 TorchScript。有关 TorchScript 注解的更多详细信息，请参阅 此处。
使用 TorchScript 而非 tracing 的一个特别原因是 tracing 无法捕获 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);
}
现在，就像之前一样，我们可以在 script 代码中使用自定义算子，就像使用任何其他函数一样：
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 导入包含你的算子（或多个算子）的库，然后在你的 traced 或 scripted TorchScript 代码中像调用任何其他 torch 算子一样调用你的自定义算子。
在 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 共享库链接起来。
注意
上面的示例中隐藏了一个关键细节：在 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++ 应用程序中进行推理工作负载。现在你已经准备好使用 C++ 算子扩展你的 TorchScript 模型，这些算子可以与第三方 C++ 库交互，编写自定义的高性能 CUDA 内核，或实现任何其他需要 Python、TorchScript 和 C++ 平滑融合的使用场景。
一如既往，如果你遇到任何问题或有疑问，可以使用我们的论坛或GitHub Issues与我们联系。此外，我们的常见问题 (FAQ) 页面可能包含有用的信息。
附录 A：更多构建自定义算子的方法¶
“构建自定义算子”一节解释了如何使用 CMake 将自定义算子构建成共享库。本附录概述了另外两种编译方法。这两种方法都使用 Python 作为编译过程的“驱动器”或“接口”。此外，这两种方法都重用了 PyTorch 为现有基础设施提供的*C++ 扩展*。C++ 扩展是 TorchScript 自定义算子的香草版（eager 模式）对应物，它们依赖于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>
