C++部署pytorch训练的模型如何设置batch_size

最近在用libtorch部署pytorch训练好的模型，由于是第一次使用libtorch，中间遇到不少问题，这次主要记录的是模型推理时如何设置batch_size>1，即一次推理多张图片，网上的教程大多都是推理1张图片(batch_size=1)，同时多张图片推理完之后，如何将输出的tensor转换为cv::Mat类型也是花费了不少时间解决，也记录一下。

一、数据读取、预处理

这部分是数据读取、处理及模型的加载，网上教程很多，这里只做记录，不多做解释。

std::string input_img="";
std::string model_path="";
//读取图片
cv::Mat img = cv::imread(input_img);
cv::cvtColor(img,img,cv::COLOR_BGR2RGB);

//加载模型
torch::NoGradGuard no_grad;
torch::jit::script::Module module=torch::jit::load(model_path);
module.to(device);
module.eval();

//mat转tensor
torch::Tensor img_tensor = torch::from_blob(img.data,{1,img.rows,img.cols,3},torch::kByte);
img_tensor = img_tensor.permute({0,3,1,2});
img_tensor = img_tensor.toType(torch::kFloat);

//归一化
img_tensor = img_tensor.div(255);
img_tensor[0][0]=img_tensor[0][0].sub_(0.485).div_(0.229);
img_tensor[0][1]=img_tensor[0][1].sub_(0.456).div_(0.224);
img_tensor[0][2]=img_tensor[0][2].sub_(0.406).div_(0.225);
img_tensor = img_tensor.to(torch::kCUDA);

二、输入batch_size=1时的推理

batch_size=1时的推理过程按照官方给的步骤操作即可。

1、batch_size=1推理

torch::Tensor output_tensor=module.forward({img_tensor}).toTensor();
output_tensor=output_tensor.squeeze(1); // [B,C,H,W]->[B,H,W]
output_tensor=output_tensor.round().to(torch::KU8);

2、转换输出的结果

cv::Mat result_mat(img.rows,img.cols,CV_8U,output_tensor.data_ptr());

三、输入batch_size>1时的推理

这次主要是记录怎么同时推理多张图片，解决方法来自Stack Overflow上的问题How to give a batch of frames to the model in pytorch c++ api?

1、batch_size=5推理

为了模拟多张图片，我这里将一个img_tensor重复输入了5次。

std::vector<torch::Tensor> tensor_vec;
tensor_vec.push_back(img_tensor);
tensor_vec.push_back(img_tensor);
tensor_vec.push_back(img_tensor);
tensor_vec.push_back(img_tensor);
tensor_vec.push_back(img_tensor);

torch::TensorList tensor_list{tensor_vec};
torch::Tensor batch_of_tensors =torch::cat(tensor_list);

torch::Tensor output_tensor=module.forward({batch_of_tensors}).toTensor();
output_tensor=output_tensor.squeeze(1); // [B,C,H,W]->[B,H,W]
output_tensor=output_tensor.round().to(torch::KU8);

2、推理结果转换为cv::Mat

解决方法来自Pytorch官方社区的一个提问Libtorch C++ convert a Tensor to cv:Mat (single channel，这个问题下的一个回答给了解决思路。

主要就是分channel来转换。

std::Vector<cv::Mat> result_mat_vec;
for(size_t k=0; k<5;K++)
{
	cv::Mat result_mat_temp(512,512,CV_8U,output_tensor[k].data_ptr());
	cv::Mat result_mat = result_mat_temp.clone();
	result_mat_vec.push_back(result_mat);
}