在C++中加载Pytorch模型

在C++中加载Pytorch模型

下面几段将概述PyTorch提供的从现有Python模型到序列化表示的路径，该序列化表示完全可以从c++加载和执行，不依赖于Python。

第一步：将Pytorch模型转换为Torch脚本

Torch Script 是Pytorch模型的一种表示，可以被Torch Script编译器所理解，编译和序列化。在多数情况下，将模型转化为Torch Script只需要很少的工作。
将PyTorch模型转换为Torch脚本有两种方法。第一种方法称为追踪（tracing），这是一种机制，通过使用示例输入对模型进行一次评估，并记录这些输入在模型中的流动，从而捕获模型的结构。这适用于有限使用控制流的模型。第二种方法是向模型添加显式注释，通知Torch脚本编译器它可以直接解析和编译您的模型代码，这取决于Torch脚本语言施加的约束。

通过追踪转换

要通过跟踪将PyTorch模型转换为Torch脚本，您必须将模型的一个实例和一个示例输入传递给 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脚本编写模型，并相应地对模型进行注释。例如，假设您有以下朴素的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

由于该模块的正向方法使用依赖于输入的控制流，不适合追踪。相反，我们可以通过从torch.jit子类化ScriptModule来将其转换为ScriptModule。ScriptModule并添加@torch.jit.script_method注释到模型的正向方法:

import torch

class MyModule(torch.jit.ScriptModule):
    def __init__(self, N, M):
        super(MyModule, self).__init__()
        self.weight = torch.nn.Parameter(torch.rand(N, M))

    @torch.jit.script_method
    def forward(self, input):
        if bool(input.sum() > 0):
          output = self.weight.mv(input)
        else:
          output = self.weight + input
        return output

my_script_module = MyModule(2, 3)

现在，创建一个新的MyModule对象将直接生成ScriptModule的一个实例，该实例已准备好进行序列化。

第二步：将脚本模块序列化到文件中

现在，创建一个新的MyModule对象将直接生成ScriptModule的一个实例，该实例已准备好进行序列化。一旦掌握了ScriptModule(通过跟踪或注释PyTorch模型)，就可以将其序列化到文件中。稍后，您将能够用c++从该文件加载模块，并在不依赖于Python的情况下执行它。假设我们想序列化前面跟踪示例中显示的ResNet18模型。要执行此序列化，只需调用模块上的save并传递一个文件名:

traced_script_module.save("model.pt")

这将在工作目录中生成model.pt文件。我们现在已经正式离开了Python的领域，并准备跨越到c++的领域。

第三步：用c++加载脚本模块

要用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;
  }

  // Deserialize the ScriptModule from a file using torch::jit::load().
  std::shared_ptr<torch::jit::script::Module> module = torch::jit::load(argv[1]);

  assert(module != nullptr);
  std::cout << "ok\n";
}

依赖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 11)

构建示例应用程序所需的最后一件事是LibTorch发行版。您可以从PyTorch网站的下载页面获取最新的稳定版本。如果你下载并解压最新的档案文件，你应该会收到一个目录结构如下的文件夹:

libtorch/
  bin/
  include/
  lib/
  share/

• lib/文件夹包含必须链接的共享库，
• include/文件夹包含程序需要包含的头文件，
• share/文件夹包含必要的CMake配置，以启用上面的简单find_package(Torch)命令。

最后一步是构建应用程序。为此，假设我们的示例目录如下所示:

example-app/
  CMakeLists.txt
  example-app.cpp

我们现在可以运行以下命令来构建应用程序内的例子-app/文件夹:

mkdir build
cd build
cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
make

其中/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模型的路径提供给最终的示例-app二进制文件，应该会得到一个友好的“ok”:

root@4b5a67132e81:/example-app/build# ./example-app model.pt
ok

第四步：在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()，它等价于torch.ones。然后运行脚本::Module的forward方法，将创建的输入向量传递给它。作为返回，我们得到一个新的IValue，我们通过调用toTensor()将其转换为一个张量。

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

第五步：获取帮助和浏览API

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

当然，有许多概念我们没有涉及。例如，您可能想要使用一个用c++或CUDA实现的自定义操作符来扩展ScriptModule，并在纯c++生产环境中加载的ScriptModule中执行这个自定义操作符。好消息是:这是可能的，并且得到了很好的支持!现在，您可以查看这个文件夹中的示例，稍后我们将提供一个教程。现时，以下连结可能会有帮助:

• The Torch Script reference: https://pytorch.org/docs/master/jit.html
• The PyTorch C++ API documentation: https://pytorch.org/cppdocs/
• The PyTorch Python API documentation: https://pytorch.org/docs/

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