在 C++ 中加载 TorchScript 模型¶

创建于：2018 年 9 月 14 日 | 最后更新：2024 年 12 月 02 日 | 最后验证：2024 年 11 月 05 日

警告

TorchScript 不再处于积极开发阶段。

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

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

PyTorch 模型从 Python 到 C++ 的旅程由 Torch Script 启用，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)

现在可以像评估常规 PyTorch 模块一样评估追踪的 ScriptModule

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：将您的 Script Module 序列化到文件¶

一旦您手中有了 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++ 中加载您的 Script Module¶

要在 C++ 中加载序列化的 PyTorch 模型，您的应用程序必须依赖 PyTorch C++ API——也称为 LibTorch。LibTorch 发行版包含共享库、头文件和 CMake 构建配置文件的集合。虽然 CMake 不是依赖 LibTorch 的必要条件，但它是推荐的方法，并且将在未来得到良好支持。在本教程中，我们将使用 CMake 和 LibTorch 构建一个最小的 C++ 应用程序，该应用程序仅加载和执行序列化的 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

如果我们提供我们之前创建的追踪 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++ 中执行 Script Module¶

在 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 的更多信息，请参阅 https://pytorch.ac.cn/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 内存中的新张量。

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

本教程有望使您对 PyTorch 模型从 Python 到 C++ 的路径有一个大致的了解。借助本教程中描述的概念，您应该能够从普通的“eager”PyTorch 模型，到 Python 中编译的 ScriptModule，到磁盘上的序列化文件，再到——为了完成循环——C++ 中可执行的 script::Module。

当然，我们还有很多概念没有介绍。例如，您可能会发现自己想使用 C++ 或 CUDA 中实现的自定义运算符扩展您的 ScriptModule，并在纯 C++ 生产环境中加载的 ScriptModule 中执行此自定义运算符。好消息是：这是可能的，并且得到了很好的支持！目前，您可以浏览此文件夹以获取示例，我们稍后将推出教程。在此期间，以下链接可能普遍有帮助

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

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