PyTorch 2.2 中文官方教程（十二）（1）https://developer.aliyun.com/article/1482557
使用 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..，其中是你的运算符名称的命名空间部分，是你的运算符的函数名称。对于我们上面编写的运算符，命名空间是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 跟踪中，就像它是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 自定义运算符集成到跟踪的 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 模型的教程，接下来的几段将基于此构建。
简而言之，即使从文件中反序列化并在 C++中运行，自定义运算符也可以像常规的torch运算符一样执行。这唯一的要求是将我们之前构建的自定义运算符共享库与我们在其中执行模型的 C++应用程序链接起来。在 Python 中，这只需简单调用torch.ops.load_library。在 C++中，您需要将共享库与您正在使用的任何构建系统中的主应用程序链接起来。以下示例将使用 CMake 展示这一点。
注意
从技术上讲，您也可以在运行时以与我们在 Python 中所做的方式相同的方式动态加载共享库到您的 C++应用程序中。在 Linux 上，您可以使用 dlopen 来做到这一点。其他平台上也存在等价物。
在上面链接的 C++执行教程的基础上，让我们从一个最小的 C++应用程序开始，该应用程序位于一个不同的文件夹中的main.cpp文件中，加载并执行一个序列化的 TorchScript 模型：
#include <torch/script.h> // One-stop header. #include <iostream> #include <memory> int main(int argc, const char* argv[]) { if (argc != 2) { std::cerr << "usage: example-app <path-to-exported-script-module>\n"; return -1; } // Deserialize the ScriptModule from a file using torch::jit::load(). torch::jit::script::Module module = torch::jit::load(argv[1]); std::vector<torch::jit::IValue> inputs; inputs.push_back(torch::randn({4, 8})); inputs.push_back(torch::randn({8, 5})); torch::Tensor output = module.forward(std::move(inputs)).toTensor(); std::cout << output << std::endl; }
还有一个小的CMakeLists.txt文件：
cmake_minimum_required(VERSION 3.1 FATAL_ERROR) project(example_app) find_package(Torch REQUIRED) add_executable(example_app main.cpp) target_link_libraries(example_app "${TORCH_LIBRARIES}") target_compile_features(example_app PRIVATE cxx_range_for)
在这一点上，我们应该能够构建应用程序：
$ mkdir build $ cd build $ cmake -DCMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" .. -- The C compiler identification is GNU 5.4.0 -- The CXX compiler identification is GNU 5.4.0 -- Check for working C compiler: /usr/bin/cc -- Check for working C compiler: /usr/bin/cc -- works -- Detecting C compiler ABI info -- Detecting C compiler ABI info - done -- Detecting C compile features -- Detecting C compile features - done -- Check for working CXX compiler: /usr/bin/c++ -- Check for working CXX compiler: /usr/bin/c++ -- works -- Detecting CXX compiler ABI info -- Detecting CXX compiler ABI info - done -- Detecting CXX compile features -- Detecting CXX compile features - done -- Looking for pthread.h -- Looking for pthread.h - found -- Looking for pthread_create -- Looking for pthread_create - not found -- Looking for pthread_create in pthreads -- Looking for pthread_create in pthreads - not found -- Looking for pthread_create in pthread -- Looking for pthread_create in pthread - found -- Found Threads: TRUE -- Found torch: /libtorch/lib/libtorch.so -- Configuring done -- Generating done -- Build files have been written to: /example_app/build $ make -j Scanning dependencies of target example_app [ 50%] Building CXX object CMakeFiles/example_app.dir/main.cpp.o [100%] Linking CXX executable example_app [100%] Built target example_app
并且在不传递模型的情况下运行它：
$ ./example_app usage: example_app <path-to-exported-script-module>
接下来，让我们序列化我们之前编写的使用我们自定义运算符的脚本函数：
torch.ops.load_library("libwarp_perspective.so") @torch.jit.script def compute(x, y): if bool(x[0][0] == 42): z = 5 else: z = 10 x = torch.ops.my_ops.warp_perspective(x, torch.eye(3)) return x.matmul(y) + z compute.save("example.pt")
最后一行将脚本函数序列化为一个名为“example.pt”的文件。如果我们将这个序列化模型传递给我们的 C++应用程序，我们可以立即运行它：
$ ./example_app example.pt terminate called after throwing an instance of 'torch::jit::script::ErrorReport' what(): Schema not found for node. File a bug report. Node: %16 : Dynamic = my_ops::warp_perspective(%0, %19)
或者也许不是。也许还不是。当然！我们还没有将自定义运算符库与我们的应用程序链接起来。让我们立即做这个，为了正确地做这件事，让我们稍微更新我们的文件组织，看起来像这样：
example_app/ CMakeLists.txt main.cpp warp_perspective/ CMakeLists.txt op.cpp
这将允许我们将warp_perspective库 CMake 目标作为我们应用程序目标的子目录。example_app文件夹中的顶层CMakeLists.txt应该如下所示：
cmake_minimum_required(VERSION 3.1 FATAL_ERROR) project(example_app) find_package(Torch REQUIRED) add_subdirectory(warp_perspective) add_executable(example_app main.cpp) target_link_libraries(example_app "${TORCH_LIBRARIES}") target_link_libraries(example_app -Wl,--no-as-needed warp_perspective) target_compile_features(example_app PRIVATE cxx_range_for)
这个基本的 CMake 配置看起来与以前很像，只是我们将warp_perspective CMake 构建添加为一个子目录。一旦它的 CMake 代码运行，我们将我们的example_app应用程序与warp_perspective共享库链接起来。
注意
上面示例中嵌入了一个关键细节：warp_perspective链接行前缀-Wl,--no-as-needed。这是必需的，因为我们实际上不会在应用程序代码中调用warp_perspective共享库中的任何函数。我们只需要TORCH_LIBRARY函数运行。不方便的是，这会让链接器混淆，并使其认为可以完全跳过与库的链接。在 Linux 上，-Wl,--no-as-needed标志强制进行链接（注意：此标志特定于 Linux！）。还有其他解决方法。最简单的方法是在您需要从主应用程序调用的运算符库中定义某个函数。这可以是在某个头文件中声明的简单函数void init();，然后在运算符库中定义为void init() { }。在主应用程序中调用此init()函数将使链接器认为这是值得链接的库。不幸的是，这超出了我们的控制范围，我们宁愿让您了解这个原因和简单的解决方法，而不是给您一些不透明的宏来放入您的代码中。
现在，由于我们现在在顶层找到了Torch包，warp_perspective子目录中的CMakeLists.txt文件可以稍微缩短一点。它应该是这样的：
find_package(OpenCV REQUIRED) add_library(warp_perspective SHARED op.cpp) target_compile_features(warp_perspective PRIVATE cxx_range_for) target_link_libraries(warp_perspective PRIVATE "${TORCH_LIBRARIES}") target_link_libraries(warp_perspective PRIVATE opencv_core opencv_photo)
让我们重新构建我们的示例应用程序，它还将链接到自定义运算符库。在顶层example_app目录中：
$ mkdir build $ cd build $ cmake -DCMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" .. -- The C compiler identification is GNU 5.4.0 -- The CXX compiler identification is GNU 5.4.0 -- Check for working C compiler: /usr/bin/cc -- Check for working C compiler: /usr/bin/cc -- works -- Detecting C compiler ABI info -- Detecting C compiler ABI info - done -- Detecting C compile features -- Detecting C compile features - done -- Check for working CXX compiler: /usr/bin/c++ -- Check for working CXX compiler: /usr/bin/c++ -- works -- Detecting CXX compiler ABI info -- Detecting CXX compiler ABI info - done -- Detecting CXX compile features -- Detecting CXX compile features - done -- Looking for pthread.h -- Looking for pthread.h - found -- Looking for pthread_create -- Looking for pthread_create - not found -- Looking for pthread_create in pthreads -- Looking for pthread_create in pthreads - not found -- Looking for pthread_create in pthread -- Looking for pthread_create in pthread - found -- Found Threads: TRUE -- Found torch: /libtorch/lib/libtorch.so -- Configuring done -- Generating done -- Build files have been written to: /warp_perspective/example_app/build $ make -j Scanning dependencies of target warp_perspective [ 25%] Building CXX object warp_perspective/CMakeFiles/warp_perspective.dir/op.cpp.o [ 50%] Linking CXX shared library libwarp_perspective.so [ 50%] Built target warp_perspective Scanning dependencies of target example_app [ 75%] Building CXX object CMakeFiles/example_app.dir/main.cpp.o [100%] Linking CXX executable example_app [100%] Built target example_app • 39
如果我们现在运行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} ]
成功！您现在已经准备好进行推理了。
流畅的 Python 第二版（GPT 重译）（九）（3）https://developer.aliyun.com/article/1482559
