使用自定义 C++ 运算符扩展 TorchScript¶
警告
本教程已在 PyTorch 2.4 中弃用。请参阅 PyTorch 自定义运算符，获取有关 PyTorch 自定义运算符的最新指南。
PyTorch 1.0 版本为 PyTorch 引入了一种新的编程模型，称为 TorchScript。TorchScript 是 Python 编程语言的一个子集，可以被 TorchScript 编译器解析、编译和优化。此外，编译后的 TorchScript 模型可以选择序列化为磁盘文件格式，您可以随后从纯 C++（以及 Python）中加载并运行该格式进行推理。
TorchScript 支持 torch 包提供的很大一部分操作，允许您纯粹地将许多种复杂模型表示为来自 PyTorch 的“标准库”的一系列张量操作。但是，在某些情况下，您可能需要使用自定义 C++ 或 CUDA 函数来扩展 TorchScript。虽然我们建议您仅在您的想法无法（以足够高的效率）表示为简单的 Python 函数时才使用此选项，但我们确实提供了一个非常友好和简单的界面，用于使用 ATen（PyTorch 的高性能 C++ 张量库）定义自定义 C++ 和 CUDA 内核。绑定到 TorchScript 后，您可以将这些自定义内核（或“操作”）嵌入到您的 TorchScript 模型中，并在 Python 和它们的序列化形式中直接在 C++ 中执行它们。
以下段落将以编写一个调用OpenCV（一个用 C++ 编写的计算机视觉库）的 TorchScript 自定义操作为例。我们将讨论如何在 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>() 方法获取指向底层数据的原始指针（就像 .data_ptr<float>() 之前用于 PyTorch 张量一样）。我们还指定张量的输出形状，我们将其硬编码为 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(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。我们进一步将一些示例输入传递给 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 模型。如果您还没有这样做，请阅读 在 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>
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++ 扩展*，它们是 TorchScript 自定义操作符的普通（急切）PyTorch 等效项，依赖于 pybind11 来“显式”绑定从 C++ 到 Python 的函数。
第一种方法使用 C++ 扩展的 便捷的即时 (JIT) 编译接口 来在你的 PyTorch 脚本首次运行时在后台编译你的代码。第二种方法依赖于久负盛名的 setuptools 包，并涉及编写一个单独的 setup.py 文件。这允许更高级的配置以及与其他基于 setuptools 的项目集成。我们将在下面详细探讨这两种方法。
使用 JIT 编译构建¶
PyTorch C++ 扩展工具包提供的 JIT 编译功能允许将自定义操作符的编译直接嵌入到你的 Python 代码中，例如在你的训练脚本的开头。
注意
这里的“JIT 编译”与在 TorchScript 编译器中优化程序时发生的 JIT 编译无关。它仅仅意味着你的自定义操作符 C++ 代码将在你第一次导入它时在你系统 /tmp 目录下的文件夹中进行编译，就好像你之前自己编译过一样。
此 JIT 编译功能有两种形式。在第一个中，你仍然将操作符实现保存在一个单独的文件 (op.cpp) 中，然后使用 torch.utils.cpp_extension.load() 来编译你的扩展。通常，此函数将返回公开你的 C++ 扩展的 Python 模块。但是，由于我们没有将自定义操作符编译成它自己的 Python 模块，我们只希望编译一个普通的共享库。幸运的是，torch.utils.cpp_extension.load() 具有一个参数 is_python_module，我们可以将其设置为 False 来表示我们只对构建共享库而不是 Python 模块感兴趣。然后 torch.utils.cpp_extension.load() 将编译并加载共享库到当前进程中，就像 torch.ops.load_library 之前所做的那样
import torch.utils.cpp_extension
torch.utils.cpp_extension.load(
name="warp_perspective",
sources=["op.cpp"],
extra_ldflags=["-lopencv_core", "-lopencv_imgproc"],
is_python_module=False,
verbose=True
)
print(torch.ops.my_ops.warp_perspective)
这应该大约打印出
<built-in method my_ops::warp_perspective of PyCapsule object at 0x7f3e0f840b10>
JIT 编译的第二种形式允许你将自定义 TorchScript 操作符的源代码作为字符串传递。为此，请使用 torch.utils.cpp_extension.load_inline
import torch
import torch.utils.cpp_extension
op_source = """
#include <opencv2/opencv.hpp>
#include <torch/script.h>
torch::Tensor warp_perspective(torch::Tensor image, torch::Tensor warp) {
cv::Mat image_mat(/*rows=*/image.size(0),
/*cols=*/image.size(1),
/*type=*/CV_32FC1,
/*data=*/image.data<float>());
cv::Mat warp_mat(/*rows=*/warp.size(0),
/*cols=*/warp.size(1),
/*type=*/CV_32FC1,
/*data=*/warp.data<float>());
cv::Mat output_mat;
cv::warpPerspective(image_mat, output_mat, warp_mat, /*dsize=*/{64, 64});
torch::Tensor output =
torch::from_blob(output_mat.ptr<float>(), /*sizes=*/{64, 64});
return output.clone();
}
TORCH_LIBRARY(my_ops, m) {
m.def("warp_perspective", &warp_perspective);
}
"""
torch.utils.cpp_extension.load_inline(
name="warp_perspective",
cpp_sources=op_source,
extra_ldflags=["-lopencv_core", "-lopencv_imgproc"],
is_python_module=False,
verbose=True,
)
print(torch.ops.my_ops.warp_perspective)
当然，最佳实践是，只有在你的源代码相对较短的情况下才使用 torch.utils.cpp_extension.load_inline。
请注意，如果你在 Jupyter Notebook 中使用它，你不应该多次执行包含注册的单元格，因为每次执行都会注册一个新的库并重新注册自定义操作符。如果你需要重新执行它，请事先重新启动笔记本的 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>
