使用C++调用并部署pytorch模型(一)

随笔 2019-03-23 3207 次浏览 次点赞

首先简述一下整体思路:

目的:使用C++及多线程可以加快模型预测速度

关于模型训练有两种方法,一种是直接使用C++编写训练代码,可以做到搭建完整的网络模型,但是无法使用迁移学习,而迁移学习是目前训练样本几乎都会用到的方法,另一种是使用python代码训练好,模型,并使用JIT技术,将python模型导出为C++可调用的模型,这里具体介绍第二种。

1.训练的模型为三种鲜花的图片分类模型:通过爬虫获取 茉莉花,栀子花,鸡蛋花(都是白色,颜色相近)的图片各300张,并使用restnet161预训练模型进行迁移学习,第5轮得到model_0.890535.pth模型(具体过程参考官方示例)

2.对模型进行转化

import torch
from torchvision import models
import torch.nn as nn

model_ft = models.densenet161(pretrained=True)

num_ftrs = model_ft.classifier.in_features

model_ft.classifier = nn.Linear(num_ftrs, 3)

model_ft.load_state_dict(torch.load('model3_0.880535.pth'))

example = torch.rand(1, 3, 224, 224)

traced_script_module = torch.jit.trace(model_ft, example)

traced_script_module.save("cxx_model.pt")

得到C++调用模型cxx_model.pt

3.根据官方API和示例编写调用代码,这里使用opencv3.x(编译过程省略)读取图片,也可以使用transforms接口读取:

#include <torch/torch.h>
#include <ATen/ATen.h>
#include <torch/script.h>

#include <iostream>
#include <memory>
#include <string>
#include <chrono>

#include "opencv2/core.hpp"
#include "opencv2/highgui.hpp"

using namespace cv;
using namespace std;
using namespace chrono;

#include <experimental/filesystem>
namespace fs = std::experimental::filesystem;

int main() {

  auto start = system_clock::now();

  std::string model_path = "/home/developer/third_code/dcgan/cxx_model.pt";
  std::string data_path = "/home/lijj/long/";
  std::string test_img = "/home/lijj/long/oranger_hair_487.jpg";
  std::string test_path = "/home/lijj/long";

  std::shared_ptr<torch::jit::script::Module> module = torch::jit::load(model_path);
  //module->to(at::kCUDA);
  assert(module != nullptr);
  std::cout << "load model ok\n";

  for (const auto &p : fs::directory_iterator(test_path)){
    std::vector<torch::jit::IValue> inputs;
    std::cout << p.path() << '\n';
    std::string s = p.path();
    Mat image = imread(s);
    std::vector<int64_t> sizes = {1, 3, image.rows, image.cols};
    at::TensorOptions options(at::ScalarType::Byte);
    at::Tensor tensor_image = torch::from_blob(image.data, at::IntList(sizes), options);
    tensor_image = tensor_image.toType(at::kFloat);
    std::ifstream is(model_path, std::ifstream::binary);
    inputs.emplace_back(tensor_image);

    //inputs.push_back(torch::ones({1, 3, 224, 224}).to(at::kCUDA));
    at::Tensor result = module->forward(inputs).toTensor();
    //auto max_result = result.max(0, true);
    //std ::cout << std::get<1>(max_result);
    //auto max_index = std::get<1>(max_result).item<float>();
    //std::cout << max_index << std::endl;
    auto pred = result.argmax(1);
    std ::cout << pred << std::endl;

  }
  auto end = system_clock::now();
  auto duration = duration_cast<microseconds>(end - start);
    cout << "花费了"
         << double(duration.count()) * microseconds::period::num / microseconds::period::den
         << "秒" << endl;
    return 0;
}

项目CMakeLists.txt为

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(dcgan)
ADD_DEFINITIONS(-g)
ADD_DEFINITIONS(-std=c++17)
find_package(Torch REQUIRED)
find_package(OpenCV REQUIRED)

add_executable(dcgan dcgan.cpp)
target_link_libraries(dcgan "${TORCH_LIBRARIES}" ${OpenCV_LIBS} stdc++fs)

set_property(TARGET dcgan PROPERTY CXX_STANDARD 11)

编译指令为

  mkdir build

  cd build

  cmake -DCMAKE_PREFIX_PATH=../../libtorch ..

  make

 ./test

输出结果:

1.png

libtorch是来自 https://download.pytorch.org/libtorch/nightly/cpu/libtorch-shared-with-deps-latest.zip 官方默认提供的是cpu库文件,gpu库需要到https://pytorch.org/选择对应的版本下载或自行编译。

仍有不足之处,后续补充,由于官方文档匮乏,需要自行摸索。

参考文档:

1.https://pytorch.org/cppdocs/

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


本文由 Tony 创作,采用 署名-非商业性使用-相同方式共享 3.0,可自由转载、引用,但需署名作者且注明文章出处。

如果对您有用,您的支持将鼓励我继续创作!