在 C++ 中加载 TorchScript 模型

原文：https://pytorch.org/tutorials/advanced/cpp_export.html

顾名思义，PyTorch 的主要接口是 Python 编程语言。 尽管 Python 是许多需要动态性和易于迭代的场景的合适且首选的语言，但是在同样许多情况下，Python 的这些属性恰恰是不利的。 后者经常应用的一种环境是生产 –低延迟和严格部署要求的土地。 对于生产场景，即使仅将 C++ 绑定到 Java，Rust 或 Go 之类的另一种语言中，它也是经常选择的语言。 以下各段将概述 PyTorch 提供的从现有 Python 模型到序列化表示形式的路径，该序列化表示形式可以完全由 C++ 加载和执行，不依赖于 Python。

第 1 步：将 PyTorch 模型转换为 Torch 脚本

PyTorch 模型从 Python 到 C++ 的旅程由 Torch 脚本启用，它是 PyTorch 模型的一种表示形式，可以由 Torch 脚本编译器理解，编译和序列化。 如果您从使用原始“渴望” API 编写的现有 PyTorch 模型开始，则必须首先将模型转换为 Torch 脚本。 在最常见的情况下（如下所述），只需很少的努力。 如果您已经有了 Torch 脚本模块，则可以跳到本教程的下一部分。

有两种将 PyTorch 模型转换为 Torch 脚本的方法。 第一种称为跟踪，该机制通过使用示例输入对模型的结构进行一次评估，并记录这些输入在模型中的流量来捕获模型的结构。 这适用于有限使用控制流的模型。 第二种方法是在模型中添加显式注解，以告知 TorchScript 编译器可以根据 Torch Script 语言施加的约束直接解析和编译模型代码。

小费

您可以在官方 Torch 脚本参考中找到这两种方法的完整文档以及使用方法的进一步指导。

通过跟踪转换为 Torch 脚本

要将 PyTorch 模型通过跟踪转换为 Torch 脚本，必须将模型的实例以及示例输入传递给torch.jit.trace函数。 这将产生一个torch.jit.ScriptModule对象，并将模型评估的轨迹嵌入到模块的forward方法中：

  1. import torch
  2. import torchvision
  4. # An instance of your model.
  5. model = torchvision.models.resnet18()
  7. # An example input you would normally provide to your model's forward() method.
  8. example = torch.rand(1, 3, 224, 224)
  10. # Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
  11. traced_script_module = torch.jit.trace(model, example)

现在可以对跟踪的ScriptModule进行评估，使其与常规 PyTorch 模块相同：

  1. In[1]: output = traced_script_module(torch.ones(1, 3, 224, 224))
  2. In[2]: output[0, :5]
  3. Out[2]: tensor([-0.2698, -0.0381,  0.4023, -0.3010, -0.0448], grad_fn=<SliceBackward>)

通过注解转换为 Torch 脚本

在某些情况下，例如，如果模型采用特定形式的控制流，则可能需要直接在 Torch 脚本中编写模型并相应地注解模型。 例如，假设您具有以下原始 Pytorch 模型：

  1. import torch
  3. class MyModule(torch.nn.Module):
  4.     def __init__(self, N, M):
  5.         super(MyModule, self).__init__()
  6.         self.weight = torch.nn.Parameter(torch.rand(N, M))
  8.     def forward(self, input):
  9.         if input.sum() > 0:
  10.           output = self.weight.mv(input)
  11.         else:
  12.           output = self.weight + input
  13.         return output

因为此模块的forward方法使用取决于输入的控制流，所以它不适合跟踪。 相反，我们可以将其转换为ScriptModule。 为了将模块转换为ScriptModule，需要使用torch.jit.script编译模块，如下所示：

  1. class MyModule(torch.nn.Module):
  2.     def __init__(self, N, M):
  3.         super(MyModule, self).__init__()
  4.         self.weight = torch.nn.Parameter(torch.rand(N, M))
  6.     def forward(self, input):
  7.         if input.sum() > 0:
  8.           output = self.weight.mv(input)
  9.         else:
  10.           output = self.weight + input
  11.         return output
  13. my_module = MyModule(10,20)
  14. sm = torch.jit.script(my_module)

如果您需要在nn.Module中排除某些方法，因为它们使用的是 TorchScript 不支持的 Python 函数，则可以使用@torch.jit.ignore来注解这些方法

my_module是已准备好进行序列化的ScriptModule的实例。

第 2 步：将脚本模块序列化为文件

跟踪或注解 PyTorch 模型后，一旦有了ScriptModule，就可以将其序列化为文件了。 稍后，您将能够使用 C++ 从此文件加载模块并执行它，而无需依赖 Python。 假设我们要序列化先前在跟踪示例中显示的ResNet18模型。 要执行此序列化，只需在模块上调用save并为其传递文件名：

  1. traced_script_module.save("traced_resnet_model.pt")

这将在您的工作目录中生成一个traced_resnet_model.pt文件。 如果您还想序列化my_module，请致电my_module.save("my_module_model.pt")。我们现在已经正式离开 Python 领域，并准备跨入 C++ 领域。

第 3 步：在 C++ 中加载脚本模块

要在 C++ 中加载序列化的 PyTorch 模型，您的应用必须依赖于 PyTorch C++ API –也称为 LibTorch 。 LibTorch 发行版包含共享库，头文件和 CMake 构建配置文件的集合。 虽然 CMake 不是依赖 LibTorch 的要求，但它是推荐的方法，将来会得到很好的支持。 对于本教程，我们将使用 CMake 和 LibTorch 构建一个最小的 C++ 应用，该应用简单地加载并执行序列化的 PyTorch 模型。

最小的 C++ 应用

让我们从讨论加载模块的代码开始。 以下将已经做：

  1. #include <torch/script.h> // One-stop header.
  3. #include <iostream>
  4. #include <memory>
  6. int main(int argc, const char* argv[]) {
  7.   if (argc != 2) {
  8.     std::cerr << "usage: example-app <path-to-exported-script-module>\n";
  9.     return -1;
  10.   }
  12.   torch::jit::script::Module module;
  13.   try {
  14.     // Deserialize the ScriptModule from a file using torch::jit::load().
  15.     module = torch::jit::load(argv[1]);
  16.   }
  17.   catch (const c10::Error& e) {
  18.     std::cerr << "error loading the model\n";
  19.     return -1;
  20.   }
  22.   std::cout << "ok\n";
  23. }

<torch/script.h>标头包含了运行示例所需的 LibTorch 库中的所有相关包含。 我们的应用接受序列化的 PyTorch ScriptModule的文件路径作为其唯一的命令行参数，然后继续使用torch::jit::load()函数对该模块进行反序列化，该函数将该文件路径作为输入。 作为回报，我们收到一个torch::jit::script::Module对象。 我们将稍后讨论如何执行它。

依赖 LibTorch 并构建应用

假设我们将以上代码存储到名为example-app.cpp的文件中。 最小的CMakeLists.txt构建起来看起来很简单：

  1. cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
  2. project(custom_ops)
  4. find_package(Torch REQUIRED)
  6. add_executable(example-app example-app.cpp)
  7. target_link_libraries(example-app "${TORCH_LIBRARIES}")
  8. set_property(TARGET example-app PROPERTY CXX_STANDARD 14)

建立示例应用的最后一件事是 LibTorch 发行版。 您可以随时从 PyTorch 网站上的下载页面获取最新的稳定版本。 如果下载并解压缩最新的归档文件，则应该收到具有以下目录结构的文件夹：

  1. libtorch/
  2.   bin/
  3.   include/
  4.   lib/
  5.   share/
  • lib/文件夹包含您必须链接的共享库，
  • include/文件夹包含程序需要包含的头文件，
  • share/文件夹包含必要的 CMake 配置，以启用上面的简单find_package(Torch)命令。

小费

在 Windows 上，调试和发行版本不兼容 ABI。 如果计划以调试模式构建项目，请尝试使用 LibTorch 的调试版本。 另外，请确保在下面的cmake --build .行中指定正确的配置。

最后一步是构建应用。 为此，假定示例目录的布局如下：

  1. example-app/
  2.   CMakeLists.txt
  3.   example-app.cpp

现在，我们可以运行以下命令从example-app/文件夹中构建应用：

  1. mkdir build
  2. cd build
  3. cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
  4. cmake --build . --config Release

其中/path/to/libtorch应该是解压缩的 LibTorch 发行版的完整路径。 如果一切顺利，它将看起来像这样：

  1. root@4b5a67132e81:/example-app# mkdir build
  2. root@4b5a67132e81:/example-app# cd build
  3. root@4b5a67132e81:/example-app/build# cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
  4. -- The C compiler identification is GNU 5.4.0
  5. -- The CXX compiler identification is GNU 5.4.0
  6. -- Check for working C compiler: /usr/bin/cc
  7. -- Check for working C compiler: /usr/bin/cc -- works
  8. -- Detecting C compiler ABI info
  9. -- Detecting C compiler ABI info - done
  10. -- Detecting C compile features
  11. -- Detecting C compile features - done
  12. -- Check for working CXX compiler: /usr/bin/c++
  13. -- Check for working CXX compiler: /usr/bin/c++ -- works
  14. -- Detecting CXX compiler ABI info
  15. -- Detecting CXX compiler ABI info - done
  16. -- Detecting CXX compile features
  17. -- Detecting CXX compile features - done
  18. -- Looking for pthread.h
  19. -- Looking for pthread.h - found
  20. -- Looking for pthread_create
  21. -- Looking for pthread_create - not found
  22. -- Looking for pthread_create in pthreads
  23. -- Looking for pthread_create in pthreads - not found
  24. -- Looking for pthread_create in pthread
  25. -- Looking for pthread_create in pthread - found
  26. -- Found Threads: TRUE
  27. -- Configuring done
  28. -- Generating done
  29. -- Build files have been written to: /example-app/build
  30. root@4b5a67132e81:/example-app/build# make
  31. Scanning dependencies of target example-app
  32. [ 50%] Building CXX object CMakeFiles/example-app.dir/example-app.cpp.o
  33. [100%] Linking CXX executable example-app
  34. [100%] Built target example-app

如果我们提供到先前创建的跟踪ResNet18模型traced_resnet_model.pt到生成的example-app二进制文件的路径，则应该以友好的“确定”来回报。 请注意，如果尝试使用my_module_model.pt运行此示例，则会收到一条错误消息，提示您输入的形状不兼容。 my_module_model.pt期望使用 1D 而不是 4D。

  1. root@4b5a67132e81:/example-app/build# ./example-app <path_to_model>/traced_resnet_model.pt
  2. ok

步骤 4：在 C++ 中执行脚本模块

在用 C++ 成功加载序列化的ResNet18之后，我们现在离执行它仅几行代码了！ 让我们将这些行添加到 C++ 应用的main()函数中：

  1. // Create a vector of inputs.
  2. std::vector<torch::jit::IValue> inputs;
  3. inputs.push_back(torch::ones({1, 3, 224, 224}));
  5. // Execute the model and turn its output into a tensor.
  6. at::Tensor output = module.forward(inputs).toTensor();
  7. 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 C++ API 提供了与 Python API 几乎相同的功能，使您可以像在 Python 中一样进一步操纵和处理张量。

在最后一行，我们打印输出的前五个条目。 由于在本教程前面的部分中，我们为 Python 中的模型提供了相同的输入，因此理想情况下，我们应该看到相同的输出。 让我们通过重新编译我们的应用并以相同的序列化模型运行它来进行尝试：

  1. root@4b5a67132e81:/example-app/build# make
  2. Scanning dependencies of target example-app
  3. [ 50%] Building CXX object CMakeFiles/example-app.dir/example-app.cpp.o
  4. [100%] Linking CXX executable example-app
  5. [100%] Built target example-app
  6. root@4b5a67132e81:/example-app/build# ./example-app traced_resnet_model.pt
  7. -0.2698 -0.0381  0.4023 -0.3010 -0.0448
  8. [ Variable[CPUFloatType]{1,5} ]

作为参考，Python 以前的输出为：

  1. tensor([-0.2698, -0.0381,  0.4023, -0.3010, -0.0448], grad_fn=<SliceBackward>)

看起来很不错！

小费

要将模型移至 GPU 内存，可以编写model.to(at::kCUDA);。 通过调用tensor.to(at::kCUDA)来确保模型的输入也位于 CUDA 内存中，这将在 CUDA 内存中返回新的张量。

第 5 步：获得帮助并探索 API

本教程有望使您对 PyTorch 模型从 Python 到 C++ 的路径有一个大致的了解。 利用本教程中介绍的概念，您应该能够从原始的“急切的” PyTorch 模型，到 Python 中的已编译ScriptModule，再到磁盘上的序列化文件，以及–结束循环–到可执行文件script::Module在 C++ 中。

当然，有许多我们没有介绍的概念。 例如，您可能会发现自己想要扩展使用 C++ 或 CUDA 实现的自定义运算符来扩展ScriptModule，并希望在纯 C++ 生产环境中加载的ScriptModule内执行该自定义运算符。 好消息是：这是可能的，并且得到了很好的支持！ 现在，您可以浏览这个文件夹作为示例，我们将很快提供一个教程。 目前，以下链接通常可能会有所帮助：

与往常一样，如果您遇到任何问题或疑问，可以使用我们的论坛或 GitHub ISSUE 进行联系。