C++部署pytorch模型例子

如题

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

有两种将PyTorch模型转换为Torch脚本的方法。第一种称为跟踪，一种机制，其中通过使用示例输入对模型的结构进行一次评估，并记录这些 输入在模型中的流量，从而捕获模型的结构。这适用于有限使用控制流的模型。第二种方法是在模型中添加显式批注，以告知Torch Script编 译器可以根据Torch Script语言施加的约束直接解析和编译模型代码。
这里以第二种为例：

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

my_module = MyModule(10,20)
sm = torch.jit.script(my_module)
sm.save('model.pt')

保存上述代码为example-app.py, 运行后生成model.pt文件

步骤2：在C++中加载脚本模块

• 编写一个example-app.cpp文件，内容如下

  #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 {
      // 使用以下命令从文件中反序列化脚本模块: 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对象。

• 编写 CMakeLists.txt，内容如下：

  cmake_minimum_required(VERSION 3.15 FATAL_ERROR)
  project(example-app)
  
  set(CMAKE_CXX_STANDARD 14)
  
  // 指定libtorch的安装目录，这里也可以使用cmake的时候指定，cmake参数:-DCMAKE_PREFIX_PATH
  set(Torch_DIR /share/apps/libtorch/libtorch-gpu/share/cmake/Torch)
  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)

• 构建应用程序

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

• 最后执行程序

  ./example-app model.pt

  正常输出为ok

参考：https://pytorch.panchuang.net/EigthSection/torchScript_in_C++/