在C++中加载TorchScript模型

创建日期: 2018年9月14日 | 最后更新日期: 2024年12月2日 | 最后验证日期: 2024年11月5日

警告

TorchScript 已经不再处于活跃开发阶段。

正如其名，PyTorch的主要接口是Python编程语言。虽然在需要动态性和易于迭代的许多场景中，Python是一种合适且首选的语言，但同样也有许多情况下这些特性并不利于使用。后者常常适用于生产环境——低延迟和严格的部署要求之地。对于生产场景，C++通常是最受欢迎的选择，即使只是为了将其绑定到另一种语言如Java、Rust或Go。接下来的段落将概述PyTorch提供的从现有Python模型转换为纯C++加载和执行的序列化表示的方法，无需依赖Python。

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

一个PyTorch模型从Python迁移到C++的过程是由Torch Script实现的，Torch Script是一种可以被Torch Script编译器理解、编译和序列化的PyTorch模型表示。如果你是从现有的使用原生“懒惰”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: 将您的脚本模块序列化到文件

一旦你手上有通过跟踪或标注PyTorch模型得到的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++中加载您的脚本模块

要将您的序列化PyTorch模型加载到C++中，您的应用程序必须依赖于 PyTorch C++ API——也称为LibTorch。LibTorch分发包包括一组共享库、头文件和CMake构建配置文件。虽然使用LibTorch并不强制要求CMake，但它是推荐的方法，并且将来会得到良好的支持。对于本教程，我们将使用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";
}

The <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/

• The lib/ 文件夹包含你需要链接的共享库。

• The include/ 文件包含你的程序需要包含的头文件，

• The 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

where /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的路径提供给生成的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的相关信息，请参阅其文档：https://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内存中的张量。

第5步: 获取帮助和探索API

本教程希望让您对一个PyTorch模型从Python到C++的路径有了基本的理解。掌握了本教程中描述的概念，您应该能够从一个基础的“懒加载”PyTorch模型，到在Python中编译成ScriptModule，再到磁盘上的序列化文件，并最终——完成循环——生成一个可执行的C++script::Module。

当然，还有很多概念我们没有涉及。例如，你可能会发现自己想要用C++或CUDA实现一个自定义操作符来扩展你的ScriptModule，并在纯C++生产环境中加载的ScriptModule中执行这个自定义操作符。好消息是：这是可能的，并且得到了很好的支持！目前，你可以浏览此文件夹中的示例，我们很快会跟进一个教程。在此期间，以下链接可能会对你有所帮助：

• Torch Script 参考：https://pytorch.org/docs/master/jit.html

• PyTorch C++ API 文档: https://pytorch.org/cppdocs/

• PyTorch Python API 文档：https://pytorch.org/docs/

如您遇到任何问题或有任何疑问，您可以使用我们的 论坛 或 GitHub 问题 来联系我们。