在 C++ 中加载 TorchScript 模型¶

创建时间： 2018年9月14日 |上次更新时间：2024 年 12 月 2 日 |上次验证： Nov 05， 2024

警告

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

顾名思义，PyTorch 的主要接口是 Python 程序设计语言。虽然 Python 是一种适合和首选的语言许多需要动态性和易于迭代的场景，同样有很多正是 Python 的这些属性不利的情况。一后者通常适用的环境是生产 - 土地低延迟和严格的部署要求。对于生产方案，C++ 通常是首选的语言，即使只是为了将其绑定到另一种语言 Java、Rust 或 Go 等语言。以下段落将概述路径 PyTorch 提供了从现有 Python 模型到序列化表示形式，可以完全从 C++ 加载和执行，没有对 Python 的依赖。

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

PyTorch 模型从 Python 到 C++ 的旅程由 Torch 脚本（PyTorch 的表示形式）实现模型，该模型可以被 Torch 脚本理解、编译和序列化编译器。如果您从用 vanilla “eager” API 时，您必须先将模型转换为 Torch 脚本。在在大多数常见情况下，下面将对此进行讨论，这只需要很少的努力。如果你已经有一个 Torch Script 模块，你可以跳到本节的下一节教程。

有两种方法可以将 PyTorch 模型转换为 Torch 脚本。第一个称为跟踪，其中模型的结构是通过使用示例输入对它进行一次评估，并记录这些输入通过模型。这适用于制作有限的模型使用控制流。第二种方法是向模型，通知 Torch 脚本编译器它可以直接解析和编译模型代码，但要遵守 Torch 脚本施加的约束语言。

提示

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

通过跟踪转换为 Torch 脚本¶

要通过跟踪将 PyTorch 模型转换为 Torch 脚本，您必须将实例以及函数的示例输入。这将生成一个带有跟踪的模型评估中嵌入的：torch.jit.tracetorch.jit.ScriptModuleforward

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 脚本¶

在某些情况下，例如，如果您的模型采用特定形式的 control flow 中，您可能希望直接在 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

因为这个模块的方法使用控制流，即取决于 Input，则不适合跟踪。相反，我们可以将它到 a . 为了将模块转换为，需要编译模块，如下所示：forwardScriptModuleScriptModuletorch.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)

如果你需要排除其中的某些方法，因为它们使用了 TorchScript 尚不支持的 Python 功能，你可以用nn.Module@torch.jit.ignore

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

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

一旦你手里有了一个，无论是从追踪还是注释 PyTorch 模型后，您就可以将其序列化为文件了。后来，您将能够在 C++ 中从此文件加载模块，并在没有对 Python 的任何依赖。假设我们想要序列化所示的模型在跟踪示例中的前面部分。要执行此序列化，只需在模块上调用 save 并为其传递文件名：ScriptModuleResNet18

traced_script_module.save("traced_resnet_model.pt")

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

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

要在 C++ 中加载序列化的 PyTorch 模型，您的应用程序必须依赖于 PyTorch C++ API – 也称为 LibTorch。LibTorch 发行版包含共享库、头文件和 CMake 构建的集合配置文件。虽然 CMake 不是必需的，但依赖于 LibTorch 的 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 的路径作为其唯一的命令行参数然后继续使用该函数反序列化模块，该函数将此文件路径作为输入。作为回报，我们收到一个对象。我们稍后将研究如何执行它。<torch/script.h>ScriptModuletorch::jit::load()torch::jit::script::Module

取决于 LibTorch 并构建应用程序¶

假设我们将上述代码存储到一个名为 .一个 minimal 构建它可能看起来很简单：example-app.cppCMakeLists.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/
该文件夹包含启用上述简单命令所需的 CMake 配置。share/find_package(Torch)

提示

在 Windows 上，调试和发布版本与 ABI 不兼容。如果您打算在 Debug 模式下构建您的项目，请尝试 LibTorch 的 debug 版本。此外，请确保在下面的行中指定正确的配置。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

其中应该是解压后的 LibTorch 的完整路径分配。如果一切顺利，它将如下所示：/path/to/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

如果我们提供之前创建的跟踪模型的路径对于生成的二进制文件，我们应该得到一个友好的 “好的”。请注意，如果尝试使用 you 运行此示例，则会收到一条错误消息，指出您的输入的形状不兼容。期望 1D 而不是 4D。ResNet18traced_resnet_model.ptexample-appmy_module_model.ptmy_module_model.pt

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

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

在 C++ 中成功加载了我们的序列化后，我们现在只是一个只需几行代码即可执行它！让我们将这些行添加到我们的 C++ 中应用程序的功能：ResNet18main()

// 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';

前两行设置了模型的输入。我们创建一个 vector of （类型擦除的值类型方法 accept 并返回）并添加单个输入。要创建输入张量，我们使用，相当于 C++ API 中的等效值。然后运行 's 方法，将输入向量传递给它，我们创建。作为回报，我们得到一个新的，我们将其转换为张量叫。torch::jit::IValuescript::Moduletorch::ones()torch.onesscript::ModuleforwardIValuetoTensor()

提示

要了解有关 PyTorch C++ API 等函数的更多信息，请参阅 general，请参阅其文档 https://pytorch.org/cppdocs。这 PyTorch C++ API 提供与 Python API 接近的奇偶校验功能，允许 u 进一步操作和处理张量，就像在 Python 中一样。torch::ones

在最后一行中，我们打印输出的前五个条目。由于我们在本教程前面的 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 内存，您可以编写 . 确保模型的输入也位于 CUDA 内存中通过调用，这将返回 CUDA 中的新张量记忆。model.to(at::kCUDA);tensor.to(at::kCUDA)

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

本教程希望能让您对 PyTorch 模型从 Python 到 C++ 的路径。通过本文中描述的概念教程中，您应该能够从原版的 “eager” PyTorch 模型转变为在 Python 中编译为磁盘上的序列化文件，并将 – 转换为 close the loop – 到 C++ 中的可执行文件。ScriptModulescript::Module

当然，我们没有涵盖许多概念。例如，您可能会发现自己想要使用自定义运算符扩展在 C++ 或 CUDA 中实现，并在纯 C++ 生产环境中加载的 Loaded 中执行此自定义运算符。好消息 IS：这是可能的，并且得到了很好的支持！现在，您可以浏览此文件夹例如，我们很快就会跟进一个教程。目前，以下链接通常可能会有所帮助：ScriptModuleScriptModule

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

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