在C++中加载TorchScript模型¶

Created On: Sep 14, 2018 | Last Updated: Dec 02, 2024 | Last Verified: Nov 05, 2024

警告

TorchScript 不再处于活跃开发状态。

顾名思义，PyTorch的主要界面是Python编程语言。虽然Python是满足许多需要动态性和快速迭代场景的合适且首选的语言，但同样有许多情况下，Python的这些特性是不利的。其中一种情况通常是 生产环境 – 低延迟和严格部署要求的领域。在生产场景中，C++往往是首选语言，即使只是为了将其绑定到另一种语言（如Java、Rust或Go）。以下段落将概述PyTorch提供的从现有Python模型到完全可以用C++加载和执行（无须依赖Python）的序列化表示的路径。

步骤1：将PyTorch模型转换为Torch Script¶

PyTorch模型从Python到C++转换的实现由 Torch Script 启用，Torch Script是一种可以被Torch Script编译器理解、编译和序列化的PyTorch模型表示形式。如果您从用原生PyTorch“延迟执行”API编写的现有模型开始，您必须首先将模型转换为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，即可将其序列化到文件中。稍后，您可以在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";
}

<torch/script.h> 头文件包含了运行该示例所需的所有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 17)

构建示例应用程序的最后一件事是获取LibTorch分发版。您始终可以从PyTorch官网上的下载页面获取最新稳定版本。如果您下载并解压最新存档，则应看到一个具有以下目录结构的文件夹：

libtorch/
  bin/
  include/
  lib/
  share/

“lib/” 文件夹包含需要链接的共享库。
“include/” 文件夹包含程序需要包含的头文件。
“share/” 文件夹包含所需的 CMake 配置，以支持上述简单的“find_package(Torch)”命令。

小技巧

在 Windows 上，调试版本和发布版本的构建不具有 ABI 兼容性。如果计划在调试模式下构建项目，请尝试使用 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

如果我们向生成的“example-app”二进制文件提供之前创建的转存的“ResNet18”模型“traced_resnet_model.pt”，我们应该会得到一个友好的“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”的向量（脚本模块方法接收和返回的类型擦除值类型），并添加一个输入。要创建输入张量，我们使用“torch::ones()”，这是 C++ API 中“torch.ones”的等效函数。之后我们运行脚本模块的“forward”方法，并传入我们创建的输入向量。返回值为一个新的“IValue”，我们可以通过调用“toTensor()”将其转换为张量。

小技巧

要了解更多像“torch::ones”这样的方法以及 PyTorch C++ API 的其他内容，请参阅其文档：https://pytorch.org/cppdocs。PyTorch C++ API 提供了与 Python API 类似的功能，使您能够像在 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 内存中的张量。

步骤 5：获取帮助并探索 API¶

希望本教程已经让您对 PyTorch 模型从 Python 到 C++ 的路径有了总体了解。通过本教程中描述的概念，您应该能够从一个普通的“即时”PyTorch 模型，到 Python 中编译的“ScriptModule”，再到磁盘上的序列化文件，并最终到可以执行的 C++ 中的“script::Module”。

当然，还有很多我们没有介绍的概念。例如，您可能希望通过在 C++ 或 CUDA 中实现自定义操作符，扩展您的“ScriptModule”，并在纯 C++ 生产环境中加载的“ScriptModule”中执行这个自定义操作符。好消息是：这是可行的，并且支持良好！目前，您可以探索这个文件夹：https://github.com/pytorch/pytorch/tree/master/test/custom_operator 来查看更多例子，我们稍后会跟进一个相关教程。同时，以下链接可能对您有所帮助：

Torch Script 参考文档：https://pytorch.org/docs/master/jit.html
PyTorch C++ API 文档：https://pytorch.org/cppdocs/
PyTorch Python API 文档：https://pytorch.org/docs/

如往常，如果您遇到任何问题或有疑问，可以使用论坛 https://discuss.pytorch.org/ 或 GitHub 问题 https://github.com/pytorch/pytorch/issues 与我们联系。