pytorch c++指南

之前说了这次的pytorch1.0有很多重大的更新，分布式更好用了，其实最重要的应该是c++的前端支持，部署更方便，效率更高，使得pytorch往产品化方面又迈进了坚实的一步。

在c++中加载pytorch模型

Step 1：将pytorch模型转化为Torch Script

首先说一下什么是Torch Script，Torch Script是PyTorch模型的一种表示，可以被Torch Script编译器理解，编译和序列化。用torch script把torch模型转成c++接口可读的模型有两种方式：Tracing && Annotation. tracing比Annotation简单，但只适合结构固定的网络模型，即forward中没有控制流的情况，因为Tracing只会保存运行时实际走的路径。如果forward函数中有控制流，需要用Annotation方式实现。

通过Tracing转换为Torch Script

tracing顾名思义，就是沿着数据运算的路径走一遍:

import torch
  import pretrainedmodels.models as models

  # An instance of your model.
  model = models.resnet50()

  # 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)

  traced_script_module.save("model.pt")

通过Annotation转换为Torch Script

Annotation要稍复杂，主要改三处：

1. Model由之前继承 nn.Model 改为继承 torch.jit.ScriptModule
2. forward函数前加 @torch.jit.script_method
3. 其他需要调用的函数前加 @torch.jit.script

Step 2: Script Module序列化

将ScriptModule序列化后才可以在c++中顺利的读取模型，而且在这个过程中不需要任何python依赖。序列化也很简单，只需要一行代码：

traced_script_module.save("model.pt")

这样得到的模型文件model.pt就可以用于c++环境当中了，而且不需要任何python环境

Step 3: c++中加载Script Module

下面我们介绍如何正确使用LibTorch来加载Script Module。

示例代码：

#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";
  }

代码很简单，从命令行加载模型文件，然后使用torch::jit::load加载。

然后是想对应的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到本地：

libtorch/
    bin/
    include/
    lib/
    share/
mkdir build
  cd build
  cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
  make

其中/path/to/libtorch必须是完整的路径

Step 4: 在c++中执行Script Module

示例代码：

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

编译&执行

make
  ./example-app model.pt

All done!

相关代码地址

• pretrainedmodels.pytorch
• pytorch-cplusplus