libtorch常用Mat与Tensor转换

        libtorch文档写的真不怎么好，看起来费劲，所以一般用pytoch训练模型，之后把模型保存为traces script model，再通过libtorch加载，进行决策，完成c++集成调用。这里对 常用libtorch API进行整理，也只整理常用的。

        验证环境：win10、vs2017、pytorch1.6、libtorch1.6、cuda10.2、python3.7。

1、Mat转Tensor

        1）将图像由BRG转为RGB通道；

        2）转为浮点型，量化到0~1；

        3）opencv转tensor；

        4）通道转换，(H,W,C)->(1,C,H,W)；

cv::cvtColor(frame, frame, CV_BGR2RGB); // step1. BGR->RGB
frame.convertTo(frame, CV_32FC3, 1.0f / 255.0f);  // step2. 归一化
// step3. 转tensor
int img_h = frame.rows;
int img_w = frame.cols;
int depth = frame.channels()
auto input_tensor = torch::from_blob(frame.data, {1, img_h, img_w, depth});
input_tensor = input_tensor.permute({0, 3, 1, 2});  // step4. 通道转换

测试代码：

        python代码逻辑：

input = torchvision.transforms.ToTensor()(img).unsqueeze(0)
output = model(input)
print("torch model out[0:5]: ", output[0, :5])

        对应c++逻辑：

cv::cvtColor(img, img, CV_BGR2RGB); // convert to RGB
img.convertTo(img, CV_32FC3, 1.0f / 255.0f); //normalization
int img_h = img.rows;
int img_w = img.cols;
int depth = img.channels();
auto input_tensor = torch::from_blob(img.data, { 1, img_h, img_w, depth }); //opencv format H*W*C
input_tensor = input_tensor.permute({ 0, 3, 1, 2 }); //pytorch format N*C*H*W

std::vector inputs;
inputs.push_back(input_tensor.to(device_type));

at::Tensor output = model.forward(inputs).toTensor();
std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

        此处也可现将图像转为tensor，之后再将tensor转为float，除以255归一化。        

cvtColor(img, img, cv::COLOR_BGR2RGB);
torch::Tensor tensor_image = torch::from_blob(img.data, { 1, img.rows, img.cols,3 }, torch::kByte);
tensor_image = tensor_image.permute({ 0,3,1,2 });
tensor_image = tensor_image.toType(torch::kFloat);

2、Tensor转Mat

        测试tensor结果或分割任务希望输出为图像时，常需将Tensor转为图像。batchsize为1时，转换分为几个步骤：

        1）equeeze去掉多余维度,(1,C,H,W)->(C,H,W)；

        2）permute执行通道顺序调整,(C,H,W)->(H,W,C)；

        3）*255，转uint8；

        4）迁移至CPU、将Tensor数据拷贝至Mat。

        5）若需要，将图像通道转回BGR。

torch::Tensor out_tensor = input_tensor;
//s1:sequeeze去掉多余维度,(1,C,H,W)->(C,H,W)；s2:permute执行通道顺序调整,(C,H,W)->(H,W,C)
out_tensor = out_tensor.squeeze().detach().permute({ 1, 2, 0 });
out_tensor = out_tensor.mul(255).clamp(0, 255).to(torch::kU8); //s3:*255，转uint8
out_tensor = out_tensor.to(torch::kCPU); //迁移至CPU
cv::Mat resultImg(img_h, img_w, CV_8UC3, out_tensor.data_ptr()); // 将Tensor数据拷贝至Mat
// cv::cvtColor(resultImg, resultImg, CV_RGB2BGR); // convert to BGR
cv::imwrite("ret.jpg", resultImg);

        tensor拷贝至Mat可分为两步(这里主要演示下memcpy函数的使用)：

cv::Mat resultImg(img_h, img_w, CV_8UC3);
std::memcpy((void *)resultImg.data, out_tensor.data_ptr(), sizeof(torch::kU8) * out_tensor.numel());

        对于输出的浮点型tensor直接转换：

cv::Mat tensor2Mat(torch::Tensor &i_tensor)
{
	int height = i_tensor.size(0), width = i_tensor.size(1);
	//i_tensor = i_tensor.to(torch::kF32);
	i_tensor = i_tensor.to(torch::kCPU);
	cv::Mat o_Mat(cv::Size(width, height), CV_32F, i_tensor.data_ptr());
	return o_Mat;
}

相关资源：

1、libtorch接口文档

Library API — PyTorch master documentation

C++ — PyTorch 1.12 documentation