c++ 使用torchscript 加载训练好的pytorch模型

1.首先官网上下载libtorch，放到当前项目下

2.将pytorch训练好的模型使用torch.jit.trace导出为.pt格式

import torch
from skimage import io, transform, color
import numpy as np
import os
import torch.nn.functional as F
import warnings
warnings.filterwarnings("ignore")
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

labels = ['cock', 'drawing', 'neutral', 'porn', 'sexy']
path = "test/n_1.jpg"
im = io.imread(path)
if im.shape[2] == 4:
    im = color.rgba2rgb(im)

im = transform.resize(im, (224, 224))
im = np.transpose(im, (2, 0, 1))
dummy_input = np.expand_dims(im, 0)
inp = torch.from_numpy(dummy_input)
inp = inp.float()
model = torch.load(
    "models/resnet50-epoch-0-accu-0.9213857428381079.pth", map_location='cpu')
traced_script_module = torch.jit.trace(model, inp)
output = model(inp)
probs = F.softmax(output).detach().numpy()[0]
pred = np.argmax(probs)

traced_script_module.save("models/traced_resnet_model.pt")

torchscript加载.pt模型

// One-stop header.
#include <torch/script.h>

// headers for opencv
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/opencv.hpp>

#include <cmath>
#include <iostream>
#include <memory>
#include <string>
#include <vector>

#define kIMAGE_SIZE 224
#define kCHANNELS 3
#define kTOP_K 1 //print top k predicted results

bool LoadImage(std::string file_name, cv::Mat &image)
{
  image = cv::imread(file_name); // CV_8UC3
  if (image.empty() || !image.data)
  {
    return false;
  }
  cv::cvtColor(image, image, CV_BGR2RGB);
  // scale image to fit
  cv::Size scale(kIMAGE_SIZE, kIMAGE_SIZE);
  cv::resize(image, image, scale);

  // convert [unsigned int] to [float]
  image.convertTo(image, CV_32FC3,1.0/255);

  return true;
}

bool LoadImageNetLabel(std::string file_name,
                       std::vector<std::string> &labels)
{
  std::ifstream ifs(file_name);
  if (!ifs)
  {
    return false;
  }
  std::string line;
  while (std::getline(ifs, line))
  {
    labels.push_back(line);
  }
  return true;
}

int main(int argc, const char *argv[])
{
  if (argc != 3)
  {
    std::cerr << "Usage:classifier <path-to-exported-script-module>  <path-to-lable-file> " << std::endl;
    return -1;
  }

  //load model
  torch::jit::script::Module module = torch::jit::load(argv[1]);
  // to GPU
  // module->to(at::kCUDA);
  std::cout << "== ResNet50 loaded!
";

  //load labels(classes names)
  std::vector<std::string> labels;
  if (LoadImageNetLabel(argv[2], labels))
  {
    std::cout << "== Label loaded! Let's try it
";
  }
  else
  {
    std::cerr << "Please check your label file path." << std::endl;
    return -1;
  }

  std::string file_name = "";
  cv::Mat image;
  while (true)
  {
    std::cout << "== Input image path: [enter q to exit]" << std::endl;
    std::cin >> file_name;
    if (file_name == "Q" || file_name == "q")
    {
      break;
    }
    if (LoadImage(file_name, image))
    {
      //read image tensor
      auto input_tensor = torch::from_blob(
          image.data, {1, kIMAGE_SIZE, kIMAGE_SIZE, kCHANNELS});
      input_tensor = input_tensor.permute({0, 3, 1, 2});
      input_tensor[0][0] = input_tensor[0][0].sub_(0.485).div_(0.229);
      input_tensor[0][1] = input_tensor[0][1].sub_(0.456).div_(0.224);
      input_tensor[0][2] = input_tensor[0][2].sub_(0.406).div_(0.225);
      // to GPU
      // input_tensor = input_tensor.to(at::kCUDA);

      torch::Tensor out_tensor = module.forward({input_tensor}).toTensor();

      auto results = out_tensor.sort(-1, true);
      auto softmaxs = std::get<0>(results)[0].softmax(0);
      auto indexs = std::get<1>(results)[0];

      for (int i = 0; i < kTOP_K; ++i)
      {
        auto idx = indexs[i].item<int>();
        std::cout << "    ============= Top-" << i + 1 << " =============" << std::endl;
        std::cout << "    Label:  " << labels[idx] << std::endl;
        std::cout << "    With Probability:  "
                  << softmaxs[i].item<float>() * 100.0f << "%" << std::endl;
      }
    }
    else
    {
      std::cout << "Can't load the image, please check your path." << std::endl;
    }
  }
}

CMakeLists.txt编译

cmake_minimum_required(VERSION 2.8)
project(predict_demo)
SET(CMAKE_CXX_FLAGS ${CMAKE_CXX_FLAGS} "-std=c++11 -O3")


set(OpenCV_DIR  /home/buyizhiyou/opencv-3.4.4/build)
find_package(OpenCV REQUIRED)
find_package(Torch REQUIRED)


# 添加头文件
include_directories( ${OpenCV_INCLUDE_DIRS} )

add_executable(resnet_demo resnet_demo.cpp)
target_link_libraries(resnet_demo ${TORCH_LIBRARIES}  ${OpenCV_LIBS})
set_property(TARGET resnet_demo PROPERTY CXX_STANDARD 11)

运行

./resnet_demo   models/traced_resnet_model.pt  labels.txt