Pytorch Tutorial 学习笔记（六）模型部署

Pytorch Tutorial 学习笔记（六）模型部署

Pytorch官方文档：https://pytorch.org/tutorials/

将pytorch模型转化到Torch Script

C++是部署的底层语言，Torch Script是连接两者的桥梁，在部署的过程中一般有两种方法：第一种称为跟踪，这是一种通过使用示例输入对其进行评估并记录这些输入在模型中的流动来捕获模型结构的机制。但是仅使用于使用了control flow的模型。第二种方法是向您的模型添加显式注释，通知 Torch 脚本编译器它可以直接解析和编译您的模型代码，受 Torch 脚本语言施加的约束。

• 通过跟踪部署

import torch
import torchvision

# 实例化模型.
model = torchvision.models.resnet18()

# 输入.
example = torch.rand(1, 3, 224, 224)

# 使用torch.jit.trace 通过跟踪生成 torch.jit.ScriptModule .
traced_script_module = torch.jit.trace(model, example)

评估跟踪效果：

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

• 通过注解部署模型

示例模型，可以看出该模型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

更改：

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 是准备好序列化的 ScriptModule 实例。

将ScriptModule序列化为文件

以上两种方法均可以将模块序列化文件，然后通过C++直接加载，无需依赖pthon。

示例保存resnet文件，

traced_script_module.save("traced_resnet_model.pt")

序列化sm文件：

sm.save("my_module_model.pt")

在C++中加载模型

要在 C++ 中加载序列化 PyTorch 模型，应用程序必须依赖于 PyTorch C++ API——也称为 LibTorch。

在官方教程中使用 CMake 和 LibTorch 构建一个最小的 C++ 应用程序，它只是加载和执行序列化的 PyTorch 模型。

#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 对象。将上述代码存储到名为 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 14)

然后在pytorch中下载最新的Libtorch，下载完成之后，得到以下目录：

libtorch/
  bin/
  include/
  lib/
  share/

lib/ 文件夹包含链接的共享库

include/ 文件夹包含程序需要包含的头文件

share/ 文件夹包含启用上述简单 find_package(Torch) 命令所需的 CMake 配置。

最后一步是构建应用程序。假设示例目录布局如下：

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的话会出现报错，原因是输入需要1D而不是4D。

root@4b5a67132e81:/example-app/build# ./example-app <path_to_model>/traced_resnet_model.pt
ok

在C++中执行模块

在 C++ 中成功加载了序列化 ResNet18 后，现在只需几行代码即可执行它！将这些行添加到我们的 C++ 应用程序的 main() 函数中：

// 创建输入向量.
std::vector<torch::jit::IValue> inputs;
inputs.push_back(torch::ones({1, 3, 224, 224}));

// 执行模型并将输出转化为tensor.
at::Tensor output = module.forward(inputs).toTensor();
std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

前两行设置了模型的输入。创建一个 torch::jit::IValue 向量（一个类型擦除的值类型 script::Module 方法接受和返回）并添加一个输入。运行 script::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 traced_resnet_model.pt
-0.2698 -0.0381  0.4023 -0.3010 -0.0448
[ Variable[CPUFloatType]{1,5} ]

输出;

tensor([-0.2698, -0.0381,  0.4023, -0.3010, -0.0448], grad_fn=<SliceBackward>)

将模型输送到GPU运行：

model.to(at::kCUDA)

保证模型在CUDA中运行：

tensor.to(at::kCUDA)

帮助文档

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