在 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 有两种方法。第一种称为*跟踪（tracing）*，这是一种机制，通过使用示例输入评估模型一次并记录这些输入在模型中的流向来捕获模型的结构。这适用于有限使用控制流的模型。第二种方法是为模型添加显式*注解（annotation）*，以告知 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：将你的 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);。通过调用 tensor.to(at::kCUDA)，确保模型的输入也位于 CUDA 内存中，这将返回一个位于 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 与我们联系。
