利用C++调用PyTorch的模型

背景

1. PyTorch的主要接口是Python语言。虽然Python是许多需要动态和易于迭代的场景的首选语言，但同样有很多情况下，Python的这些属性恰好是不利的。在生产环境中，需要保证低延迟和其它严格的要求，而对于生产场景，C++通常是首选语言，通常C++会被绑定到Java或Go当中；
2. 第一种方法是在PyTorch中直接使用C++编写PyTorch的前端，而不是通常情况下使用Python来编写PyTorch的前端，实现模型定义、训练和评估以及模型的部署；
3. 第二种方法是使用Python编写PyTorch的前端，并且实现上述功能；
4. 众所周知，Python相对于C++在不考虑执行效率的情况下具有很多优势，本文不会讨论这方面的问题。因此，如果可以使用Python编写前端，实现模型定义、训练和评估，而将模型的部署交由C++实现，则可以最大化目标，最快地获得模型以及部署高效的模型。
5. 概念转换成具体的方案，将Python在PyTorch下训练得到的模型文件转化成C++可以加载和执行的模型文件，并且自此以后不再依赖于Python；
6. PyTorch模型从Python到C ++之旅由Torch Script实现，Torch Script是PyTorch模型的一种表示，并且可以由Torch Script的编译器理解、编译和序列化。

环境

1. 安装PyTorch的版本为1.0及以上；
2. 安装C++版本的LibTorch，LibTorch发行版包含一组共享库，头文件和CMake构建配置文件；
3. 安装Intel所提供的MKL-DNN库，Caffe依赖此库，编译得到的可执行程序会依赖其中的libmklml和libiomp5动态链接库，若无此库在执行程序时会有错误产生，安装后将这两个动态链接库拷贝至LibTorch的lib目录下；

Mac OS 报错信息如下：
dyld: Library not loaded: @rpath/libmklml.dylib> dyld: Library not loaded: @rpath/libmklml.dylib
Referenced from: ******
Reason: image not found

4. 安装CMake。

将PyTorch模型转化为Torch Script的两种方法

1. 如果需要C++使用PyTorch的模型，就必须先将PyTorch模型转化为Torch Script；
2. 目前有两种方法，可以将PyTorch模型转化为Torch Script：
   • 第一种方法是tracing。该方法通过将样本输入到模型中，并对该过程进行推断从而捕获模型的结构，并记录该样本在模型中的控制流。该方法适用于模型中较少使用控制流的模型；
   • 第二种方法是向模型中添加显式的注释，使得Torch Script编译器可以直接解析和编译模型的代码，受Torch Script强加的约束。该方法适用于使用特定控制流的模型，。

利用Tracing将模型转换为Torch Script

通过tracing的方法将PyTorch的模型转换为Torch Script，则必须将模型的实例以及样本输入传递给torch.jit.trace方法。这样会生成一个torch.jit.ScriptModule对象，模型中的forward方法中用预先嵌入了模型推断的跟踪机制：

import torch
import torchvision

# An instance of your model.
model = torchvision.models.resnet18()

# model.eval()

# 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
output = traced_script_module(torch.ones(1, 3, 224, 224))

利用注释将模型转换为Torch Script

1. 在某些情况下，如模型采用特定形式的控制流（if...else...），可以使用注释的方法将模型转化为Torch Script；
2. 此模块的forward方法使用依赖于输入的控制流，因此它不适合利用tracing的方法生成Torch Script。可以通过继承torch.jit.ScriptModule并将@torch.jit.script_method标注添加到模型的forward中的方法，来将模型转换为ScriptModule。

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

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 input.sum() > 0:
          output = self.weight.mv(input)
        else:
          output = self.weight + input
        return output

my_script_module = MyModule()

3. MyModule现在直接创建一个新对象会生成一个可以进行序列化的ScriptModule实例 。

模型序列化

不论使用了上述的哪一种方法，当ScriptModule掌握了模型的Tracing或注释，就可以将其序列化为文件。稍后将能够使用C++从该文件加载模型并执行它，不需要依赖于Python。而要执行序列化，只需在模型的实例上调用save方法。

traced_script_module.save("model.pt")

---

现在正式离开Python的领域，并准备跨越到C ++领域。

---

使用C++加载脚本模块

1. 在C++中加载序列化的PyTorch模型，应用程序必须依赖于PyTorch C++ API，也称为LibTorch。LibTorch的发行版包含一组共享库，头文件和CMake构建配置文件。CMake是LibTorch推荐的方法，并且将来会得到很好的支持；

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

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

2. <torch/script.h>是编译运行上述代码所需的LibTorch头文件；
3. 接受序列化的ScriptModule文件作为唯一的命令行参数；
4. 使用torch::jit::load方法反序列化文件，该方法的参数为序列化ScriptModule的文件；
5. 反序列化文件后返回共享所有权的智能指针，其类型为torch::jit::script::Module，与Python中的torch.jit.ScriptModule相对应。

使用CMake构建应用程序

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

构建应用程序

# 目录布局

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

libtorch/
  bin/
  include/
  lib/
  share/

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

在main()函数添加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.
auto output = module->forward(inputs).toTensor();

std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

需要再次编译，创建新的可执行文件。

make

利用注释将模型转换为Torch Script的例子

# Convolutional neural network (two convolutional layers)
class ConvNet(torch.jit.ScriptModule):
    def __init__(self, num_classes=10):
        super(ConvNet, self).__init__()
        self.layer1 = torch.jit.trace(nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)), torch.rand(1, 1, 28, 28))
        self.layer2 = torch.jit.trace(nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)), torch.rand(1, 16, 14, 14))
        self.fc = nn.Linear(7*7*32, num_classes)
        
    @torch.jit.script_method
    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out

1. 父类不再是torch.nn.module，而是torch.jit.ScriptModule；
2. 使用torch.jit.trace跟踪函数，跟踪只能正确记录不依赖于数据的函数和模块，并且没有任何未跟踪的外部依赖（例如，执行输入/输出或访问全局变量）；
3. Python函数或者模块将使用输入数据执行，其参数和返回值必须是Tensor或者包含Tensor的元组，函数或者模块作为方法的第一个参数；
4. 在跟踪时将输入数据的形状传递给函数作为方法的第二个参数；

def f(x):
    return x * 2
traced_f = torch.jit.trace(f, torch.rand(1))

5. 完整的转换为Torch Script的Python代码（模型训练）：

from __future__ import print_function

import torch 
import torch.nn as nn

# Convolutional neural network (two convolutional layers)
class ConvNet(torch.jit.ScriptModule):
    def __init__(self, num_classes=10):
        super(ConvNet, self).__init__()
        self.layer1 = torch.jit.trace(nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)), torch.rand(1, 1, 28, 28))
        self.layer2 = torch.jit.trace(nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)), torch.rand(1, 16, 14, 14))
        self.fc = nn.Linear(7*7*32, num_classes)
        
    @torch.jit.script_method
    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out

traced_script_module = ConvNet()

traced_script_module.load_state_dict(torch.load('model.ckpt'))

traced_script_module.eval()

traced_script_module.save('model.pt')

with torch.no_grad():
    outputs = traced_script_module(torch.rand(1, 1, 28, 28))

print(torch.nn.functional.softmax(outputs, dim=1))

6. 完整的反序列化和执行Script Module的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";

    // Create a vector of inputs.
  std::vector<torch::jit::IValue> inputs;
  inputs.push_back(torch::rand({1, 1, 28, 28}));
  
  // Execute the model and turn its output into a tensor.
  auto output = module->forward(inputs).toTensor();
  
  std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/10) << '\n';
}

7. 结果

./example-app model.pt
ok
Columns 1 to 8-12.0483   2.1065  -2.1548 -11.6267  -6.6993  -7.9013 -12.9029  -1.5719

Columns 9 to 10-14.2974 -10.0303
[ Variable[CPUFloatType]{1,10} ]

---

官方文档