c++ pytorch部署相关操作代码 libtorch

该笔记是：
pytorch模型转成c++后，进行单图识别（推理）所写的主要代码，（这里我是做单字分类）。

头文件

#include <torch/script.h>
#include <iostream>
#include <memory>
#include <string>
#include <cmath>

// opencv 头文件
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/opencv.hpp>

加载pt模型

torch::jit::script::Module module;
module = torch::jit::load("XXX.pt");

加载图片

这里也是用的opencv读图片

cv::Mat input_image= cv::imread("XXX.png");
if (input_image.empty() || !input_image.data) // 图片异常判断
        std::cout << "read image fail" << std::endl;

图片处理

这要根据自己模型的处理方式做不同的修改，
还有其他的方式没写出来，比如缩放cv::resize等

input_image.convertTo(input_image, CV_32FC3, 1.0 / 255.0);  // 读完的图片矩阵中的int值转float

torch::Tensor tensor_image = torch::from_blob(input_image.data, {1, input_image.rows, input_image.cols,3});  // 转Tensor

tensor_image = tensor_image.permute({0,3,1,2});  //一些换维操作

//三通道都归一化，注意这里均值方差我都是0.5，自己修改
tensor_image[0][0] = tensor_image[0][0].sub_(0.5).div_(0.5);
tensor_image[0][1] = tensor_image[0][1].sub_(0.5).div_(0.5);
tensor_image[0][2] = tensor_image[0][2].sub_(0.5).div_(0.5);

前向传播

auto output = module.forward({tensor_image}).toTensor();

获得识别结果（求输出概率最高的），用torch::max

std::tuple<torch::Tensor, torch::Tensor> rec_result = torch::max(output,1);
auto rec_class= std::get<0>(rec_result);
auto rec_class_index= std::get<1>(rec_result);

std::cout << "概率值为" << log10(rec_class) << "\n\n";
std::cout << "识别结果对应index" << rec_class_index << "\n\n";

求前K个概率高的识别结果，用torch::topk

std::tuple<torch::Tensor, torch::Tensor> topk_class = torch::topk(output,10); // 第二个参数10就是前k个
auto top_10 = std::get<0>(topk_class);  // 前十个概率取e, 所以后面得到概率值要先进行log10操作
auto top_index = std::get<1>(topk_class);  // 前十个高概率字符

int k = 0; // 第k个概率高的字符，k自己改或者写个循环
std::cout << "概率值为：" << log10(top_10[0][k]) << "\n\n";  
std::cout << "对应index" << top_index[0][k] << "\n\n";