PyTorch 2.2 中文官方教程（九）（1）https://developer.aliyun.com/article/1482545
在 C++中加载 TorchScript 模型
原文：
pytorch.org/tutorials/advanced/cpp_export.html译者：飞龙
正如其名称所示，PyTorch 的主要接口是 Python 编程语言。虽然 Python 是许多需要动态性和迭代便利性的场景的合适和首选语言，但同样有许多情况下，Python 的这些特性并不理想。其中一个经常适用后者的环境是生产 - 低延迟和严格部署要求的领域。对于生产场景，C++往往是首选的语言，即使只是将其绑定到另一种语言如 Java、Rust 或 Go 中。以下段落将概述 PyTorch 提供的路径，从现有的 Python 模型到可以纯粹从 C++中加载和执行的序列化表示形式，而不依赖于 Python。
第 1 步：将您的 PyTorch 模型转换为 Torch Script
PyTorch 模型从 Python 到 C++的旅程是由Torch Script实现的，这是 PyTorch 模型的一种表示形式，可以被 Torch Script 编译器理解、编译和序列化。如果您从使用原始“eager”API 编写的现有 PyTorch 模型开始，您必须首先将您的模型转换为 Torch Script。在下面讨论的最常见情况下，这只需要很少的努力。如果您已经有了一个 Torch Script 模块，您可以跳过本教程的下一部分。
将 PyTorch 模型转换为 Torch Script 有两种方法。第一种被称为跟踪，通过使用示例输入对模型进行一次评估并记录这些输入通过模型的流程来捕获模型的结构。这适用于对控制流使用有限的模型。第二种方法是向您的模型添加显式注释，通知 Torch Script 编译器可以直接解析和编译您的模型代码，受 Torch Script 语言的约束。
提示
您可以在官方Torch Script 参考文档中找到这两种方法的完整文档，以及关于使用哪种方法的进一步指导。
通过跟踪转换为 Torch Script
要通过跟踪将 PyTorch 模型转换为 Torch Script，必须将模型实例和示例输入传递给torch.jit.trace函数。这将生成一个带有模型评估跟踪的torch.jit.ScriptModule对象，嵌入在模块的forward方法中：
import torch import torchvision # An instance of your model. model = torchvision.models.resnet18() # An example input you would normally provide to your model's forward() method. example = torch.rand(1, 3, 224, 224) # Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing. traced_script_module = torch.jit.trace(model, example)
跟踪的ScriptModule现在可以像常规 PyTorch 模块一样进行评估：
In[1]: output = traced_script_module(torch.ones(1, 3, 224, 224)) In[2]: output[0, :5] Out[2]: tensor([-0.2698, -0.0381, 0.4023, -0.3010, -0.0448], grad_fn=<SliceBackward>)
通过注释转换为 Torch Script
在某些情况下，例如如果您的模型使用特定形式的控制流，您可能希望直接在 Torch Script 中编写您的模型并相应地注释您的模型。例如，假设您有以下基本的 Pytorch 模型：
import torch class MyModule(torch.nn.Module): def __init__(self, N, M): super(MyModule, self).__init__() self.weight = torch.nn.Parameter(torch.rand(N, M)) def forward(self, input): if input.sum() > 0: output = self.weight.mv(input) else: output = self.weight + input return output
由于此模块的forward方法使用依赖于输入的控制流，因此不适合跟踪。相反，我们可以将其转换为ScriptModule。为了将模块转换为ScriptModule，需要使用torch.jit.script编译模块，如下所示：
class MyModule(torch.nn.Module): def __init__(self, N, M): super(MyModule, self).__init__() self.weight = torch.nn.Parameter(torch.rand(N, M)) def forward(self, input): if input.sum() > 0: output = self.weight.mv(input) else: output = self.weight + input return output my_module = MyModule(10,20) sm = torch.jit.script(my_module)
如果您需要排除nn.Module中的某些方法，因为它们使用 TorchScript 尚不支持的 Python 特性，您可以使用@torch.jit.ignore对其进行注释。
sm是一个准备好进行序列化的ScriptModule实例。
第 2 步：将您的脚本模块序列化到文件
一旦您手头有一个ScriptModule，无论是通过跟踪还是注释 PyTorch 模型获得的，您就可以将其序列化到文件中。稍后，您将能够在 C++中从此文件加载模块并执行它，而无需依赖 Python。假设我们想要序列化前面在跟踪示例中显示的ResNet18模型。要执行此序列化，只需在模块上调用save并传递一个文件名：
traced_script_module.save("traced_resnet_model.pt")
这将在您的工作目录中生成一个traced_resnet_model.pt文件。如果您还想序列化sm，请调用sm.save("my_module_model.pt")。我们现在正式离开了 Python 领域，准备进入 C++领域。
第 3 步：在 C++中加载您的脚本模块
在 C++中加载您序列化的 PyTorch 模型，您的应用程序必须依赖于 PyTorch C++ API - 也称为LibTorch。LibTorch 分发包括一组共享库、头文件和 CMake 构建配置文件。虽然 CMake 不是依赖于 LibTorch 的必需条件，但它是推荐的方法，并且将在未来得到很好的支持。在本教程中，我们将构建一个最小的 C++应用程序，使用 CMake 和 LibTorch 简单地加载和执行一个序列化的 PyTorch 模型。
一个最小的 C++应用程序
让我们从讨论加载模块的代码开始。以下内容已经足够：
#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; } torch::jit::script::Module module; try { // Deserialize the ScriptModule from a file using torch::jit::load(). module = torch::jit::load(argv[1]); } catch (const c10::Error& e) { std::cerr << "error loading the model\n"; return -1; } std::cout << "ok\n"; }
头文件包含了 LibTorch 库中运行示例所需的所有相关包含。我们的应用程序接受一个序列化的 PyTorch ScriptModule的文件路径作为唯一的命令行参数，然后使用torch::jit::load()函数对模块进行反序列化，该函数以此文件路径作为输入。作为返回，我们收到一个torch::jit::script::Module对象。我们将在稍后看看如何执行它。
依赖于 LibTorch 并构建应用程序
假设我们将上面的代码存储到一个名为example-app.cpp的文件中。一个用于构建它的最小CMakeLists.txt可能看起来就像这样简单：
cmake_minimum_required(VERSION 3.0 FATAL_ERROR) project(custom_ops) find_package(Torch REQUIRED) add_executable(example-app example-app.cpp) target_link_libraries(example-app "${TORCH_LIBRARIES}") set_property(TARGET example-app PROPERTY CXX_STANDARD 14)
构建示例应用程序所需的最后一件事是 LibTorch 分发。您可以随时从 PyTorch 网站的下载页面上获取最新的稳定版本。如果下载并解压最新的存档，您应该会收到一个具有以下目录结构的文件夹：
libtorch/ bin/ include/ lib/ share/
lib/文件夹包含了您必须链接的共享库，include/文件夹包含了您的程序需要包含的头文件，share/文件夹包含了必要的 CMake 配置，以启用上面简单的find_package(Torch)命令。
提示
在 Windows 上，调试版本和发布版本不兼容。如果您计划在调试模式下构建项目，请尝试使用 LibTorch 的调试版本。此外，请确保在下面的cmake --build .行中指定正确的配置。
最后一步是构建应用程序。为此，假设我们的示例目录布局如下：
example-app/ CMakeLists.txt example-app.cpp
我们现在可以运行以下命令来从example-app/文件夹中构建应用程序：
mkdir build cd build cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch .. cmake --build . --config Release
其中/path/to/libtorch应该是解压后的 LibTorch 分发的完整路径。如果一切顺利，它应该看起来像这样：
root@4b5a67132e81:/example-app# mkdir build root@4b5a67132e81:/example-app# cd build root@4b5a67132e81:/example-app/build# cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch .. -- 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 -- Configuring done -- Generating done -- Build files have been written to: /example-app/build root@4b5a67132e81:/example-app/build# make Scanning dependencies of target example-app [ 50%] Building CXX object CMakeFiles/example-app.dir/example-app.cpp.o [100%] Linking CXX executable example-app [100%] Built target example-app
如果我们将之前创建的跟踪的ResNet18模型traced_resnet_model.pt的路径提供给生成的example-app二进制文件，我们应该会得到一个友好的“ok”。请注意，如果尝试使用my_module_model.pt运行此示例，您将收到一个错误，指出您的输入形状不兼容。my_module_model.pt期望的是 1D 而不是 4D。
root@4b5a67132e81:/example-app/build# ./example-app <path_to_model>/traced_resnet_model.pt ok
第 4 步：在 C++中执行脚本模块
在 C++中成功加载我们序列化的ResNet18之后，我们现在只需再加入几行代码就可以执行它了！让我们将这些行添加到我们的 C++应用程序的main()函数中：
// Create a vector of inputs. std::vector<torch::jit::IValue> inputs; inputs.push_back(torch::ones({1, 3, 224, 224})); // Execute the model and turn its output into a tensor. at::Tensor output = module.forward(inputs).toTensor(); std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';
前两行设置了我们模型的输入。我们创建了一个torch::jit::IValue向量（一种类型擦除的值类型，script::Module方法接受和返回），并添加了一个单一的输入。为了创建输入张量，我们使用torch::ones()，相当于 C++ API 中的torch.ones。然后我们运行script::Module的forward方法，将我们创建的输入向量传递给它。作为返回，我们得到一个新的IValue，通过调用toTensor()将其转换为张量。
提示
要了解更多关于torch::ones等函数和 PyTorch C++ API 的信息，请参考其文档：pytorch.org/cppdocs。PyTorch C++ API 几乎与 Python API 具有相同的功能，允许你像在 Python 中一样进一步操作和处理张量。
在最后一行，我们打印输出的前五个条目。由于我们在本教程中之前在 Python 中向模型提供了相同的输入，我们应该理想情况下看到相同的输出。让我们尝试重新编译我们的应用程序，并使用相同的序列化模型运行它：
root@4b5a67132e81:/example-app/build# make Scanning dependencies of target example-app [ 50%] Building CXX object CMakeFiles/example-app.dir/example-app.cpp.o [100%] Linking CXX executable example-app [100%] Built target example-app root@4b5a67132e81:/example-app/build# ./example-app traced_resnet_model.pt -0.2698 -0.0381 0.4023 -0.3010 -0.0448 [ Variable[CPUFloatType]{1,5} ]
作为参考，之前在 Python 中的输出是：
tensor([-0.2698, -0.0381, 0.4023, -0.3010, -0.0448], grad_fn=<SliceBackward>)
看起来很匹配！
提示
要将模型移动到 GPU 内存，你可以写model.to(at::kCUDA);。确保模型的输入也在 CUDA 内存中，通过调用tensor.to(at::kCUDA)，这将返回一个在 CUDA 内存中的新张量。
第五步：获取帮助和探索 API
这个教程希望能让你对 PyTorch 模型从 Python 到 C++的路径有一个基本的理解。有了这个教程中描述的概念，你应该能够从一个普通的“eager” PyTorch 模型，转换为 Python 中编译的ScriptModule，再到磁盘上的序列化文件，最后到 C++中可执行的script::Module。
当然，还有许多概念我们没有涉及。例如，你可能会发现自己想要在 C++或 CUDA 中实现自定义运算符来扩展你的ScriptModule，并在纯 C++生产环境中加载这个自定义运算符并在你的ScriptModule中执行。好消息是：这是可能的，并且得到了很好的支持！目前，你可以在这里探索示例，我们将很快推出教程。在此期间，以下链接可能会有所帮助：
- Torch Script 参考：
pytorch.org/docs/master/jit.html - PyTorch C++ API 文档：
pytorch.org/cppdocs/ - PyTorch Python API 文档：
pytorch.org/docs/
如果遇到任何问题或有疑问，你可以使用我们的论坛或GitHub 问题来联系我们。
PyTorch 2.2 中文官方教程（九）（3）https://developer.aliyun.com/article/1482548
