Libtorch教程(二)：基于Libtorch分类模型推理

python 将pytorch的模型转为trace模型，这种模型包含既包含模型和权重。

import torch
import torchvision

 model = torchvision.models.resnet18(True)
 model.eval()
 ts_model = torch.jit.trace(model, torch.rand(1, 3, 224, 224))
 ts_model.save("traces_script_module.pt")

c++ libtorch推理代码如下

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

int main(int argc, char* argv[])
{
   // 1.定义设备类型
   std::cout << "cuda::is_available():" << torch::cuda::is_available() << std::endl;
   torch::DeviceType device_type = at::kCPU;
   if (torch::cuda::is_available()) {
      device_type = at::kCUDA;
   }

   //2 .加载模型
   torch::jit::script::Module model = torch::jit::load("F:\\traces_script_module.pt");
   std::cout << "load model is successed!" << std::endl;
   model.to(device_type);
   std::cout << "load model to device!" << std::endl;
   model.eval();

   //3. 读取图片
   cv::Mat img = cv::imread("F:\\bus.jpg");
   //cv::resize(img, img, cv::Size(640, 640));
   cv::resize(img, img, cv::Size(224, 224));

   cv::cvtColor(img, img, cv::COLOR_BGR2RGB);  // BGR -> RGB
   img.convertTo(img, CV_32FC3, 1.0f / 255.0f);  // normalization 1/255
   auto imgTensor = torch::from_blob(img.data, { 1, img.rows, img.cols, img.channels() }).to(device_type);
   imgTensor = imgTensor.permute({ 0, 3, 1, 2 }).contiguous();  // BHWC -> BCHW (Batch, Channel, Height, Width)

   //4.前向推理
   std::vector<torch::jit::IValue> inputs;
   inputs.emplace_back(imgTensor);
   torch::jit::IValue output = model.forward(inputs);

   auto ouputTensor = output.toTensor();
   torch::Tensor output_max = ouputTensor.argmax(1);
   int index = output_max.item().toInt();
   std::cout << "index:" << index << std::endl;
   return 0;
}