C++调用PyTorch模型

利用libtorch转化模型

Posted by kevin on May 16, 2020

preface

上次在服务器上装上了 docker 版的 OpenCV,就是为本篇文章服务的,因为最近做模型的部署,就要用 C++ 调用 pytorch 的模型,pytorch 推出了 libtorch 来进行模型的部署,本篇文章就记录一下部署的过程。

转化模型

用 pytorch 训练完的模型必须要在 python 环境下才能够被调用,既然我们要让它能够被 C++ 调用那就得转化模型,pytorch 官方提出的 TorchScript 就是用来做这事的,TorchScript 是一种从 PyTorch 代码创建可序列化和可优化模型的方法。任何 TorchScript 程序都可以从 Python 进程中保存并在没有 Python 依赖项的进程中加载。

那么具体怎么做呢,也很简单,就是用 torch.jit.trace 对模型进行追踪,拿我们训练好的 mnist 模型来做例子(部分代码未给出,只有核心代码),torch.jit.trace 函数接收模型和示例输入这两个参数,其中示例输入的尺寸要与模型的输入尺寸相同,为 (1, channels, img_width, img_height),这里 LeNet 模型,输入的图像是单通道 28*28 大小的。这样模型就被追踪了,将其序列化为 .pt 后缀的模型文件,这样我们就离开了 python 的领域,准备跨入 C++ 领域了。

import torch

model = Net().to('cpu')
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
load_checkpoint('mnist-4690.pth', model, optimizer)
example = torch.rand(1, 1, 28, 28)
traced_model = torch.jit.trace(model, example)
output = traced_model(torch.ones(1, 1, 28, 28))
print(output)
traced_model.save('traced_mnist_pytorch_model.pt')

代码的输出为:tensor([[ -1.9272, -10.6155, -0.5130, 4.6338, -7.4411, 0.6665, -8.3521, -3.3505, 16.8189, 3.7142]], grad_fn=)

下载依赖库

C++ 读图像大多用 OpenCV,我们已经编译好了,物理机编译可以参考我这篇教程,然后还要用到官方的 libtorch 库,直接去官网下载对应的版本然后解压就行了,不需要安装。

编写C++代码

接下去就利用 OpenCV 和 libtorch 对模型进行调用,注意了,测试图片的时候要将图片进行预处理,预处理步骤和在 PyTorch 中的预处理要一模一样,否则会出错。

#include <torch/script.h>
#include <iostream> 
#include <memory>  
#include <opencv2/opencv.hpp>  

using namespace std;
using namespace cv;

int main(int argc, char * argv[]){
    if (argc != 2) {
      std::cerr << "no module found !\n";
      return -1;
    }

    torch::jit::script::Module module;
    // 加载序列化的PyTorch模型
    try {
        module = torch::jit::load(argv[1]);
    }
    catch (const c10::Error& e){
        std::cerr << "error loading the module\n";
        return -1;
    }

    std::cout << "ok!\n";

    /////////////////////////////////////////////////
    
    string path = "/pytorch-deployment/assets/3.jpg";
    Mat img = imread(path), img_float;
    cvtColor(img, img, CV_BGR2RGB); //PyTorch的tensor是RGB通道,OpenCV是BGR
    bitwise_not(img, img);  // 将图像黑白反相
    vector<Mat> mv;
    split(img, mv);
    img = mv[1];    // 截取单通道
    img.convertTo(img_float, CV_32F, 1.0 / 255);    // 使图像归一化
    resize(img_float, img_float , Size(28, 28));    // 将图像resize成LeNet网络的输入大小
    auto img_tensor = torch::from_blob(img_float.data, {1, 28, 28, 1}, at::kFloat).permute({ 0,3,1,2 });    // change the channels
    auto img_var = torch::autograd::make_variable(img_tensor, false);

    // Create a vector of inputs.
    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(img_var);

    auto output = module.forward(inputs).toTensor();    // 模型前向传播得到结果
    cout << output << endl;
    auto index = output.argmax(1);  // 概率最大的onehot编码
    cout << "The predicted class is : " << index << endl;
}

编译代码

由于我是用的 docker 版本的 OpenCV,同时又要用到代码源文件和 libtorch 库,就用 docker 的 -v 将物理机的目录挂载进 docker ,并用下面命令进入 docker (不在 docker 下部署的话可以忽视这一步)

$ docker run -v /tmp/.X11-unix:/tmp/.X11-unix -e DISPLAY=$DISPLAY -p 5000:5000 -p 8889:8888 -e GRANT_SUDO=yes --user root -v ~/kevin/code/pytorch/pytorch-model-deployment:/pytorch-deployment -v ~/kevin/code/pytorch/libtorch:/libtorch -it spmallick/opencv-docker:opencv-3.4.3 /bin/bash

编译官方选用 CMake 工具,因此要写 CMakeLists.txt 文件,用 Qt Creator 构建我没有成功,以后有空再康康(主要是有代码自动补全功能)

cmake_minimum_required(VERSION 2.8.12)

set(CMAKE_INCLUDE_CURRENT_DIR ON)
set(CMAKE_AUTOMOC ON)
set(CMAKE_PREFIX_PATH ${CMAKE_PREFIX_PATH} "/libtorch")

project(pytorch-deployment)

find_package(Torch REQUIRED)
find_package(OpenCV REQUIRED)

add_executable(${PROJECT_NAME} "main.cpp")

target_link_libraries(${PROJECT_NAME} "${TORCH_LIBRARIES}")
target_link_libraries(${PROJECT_NAME} "${OpenCV_LIBS}")
set_property(TARGET ${PROJECT_NAME} PROPERTY CXX_STANDARD 14)

依旧创建一个 build 目录来装编译中间产物,在 build 目录下输入以下命令构建可执行文件

$ cmake ..
$ make -j

构建完成后直接运行即可,不过这里还需要加上模型的存放路径作为命令行参数

$ ./pytorch-deployment  /pytorch-deployment/assets/traced_mnist_pytorch_model.pt

可以得到最终的输出,本次用了一张手写数字 3 进行测试,模型最终输出也是 3,证明部署成功了,其实不出错的话,最终的结果和 PyTorch 会是一样的。

output

reference

https://pytorch.org/tutorials/advanced/cpp_export.html

c++调用pytorch模型并使用GPU进行预测_人工智能_u010397980的博客-CSDN博客