[LibTorch] Tensor 与 Mat 相互转换

Torch 框架搭建的神经网络，输入和输出的类型一般是 CUDAFloatType。如果是对图片进行处理，首先使用 OpenCV 读取图片，然后将其转换为 CUDAFloatType，经过神经网络模型进行处理后， 再把结果转为 cv::Mat，并通过 imwrite 输出图片。

单通道

神经网络输出的 tensor 的通道数为 1。

#include <torch/torch.h>
#include <torch/script.h>

/* 初始化模型 */
torch::jit::script::Module net = torch::jit::load("../models/poolnet.pt");
net.to(torch::kCUDA);
net.eval();

/* cv::Mat { h, w, 3 } -> CUDAFloatType { 1, 3, h, w } */
cv::Mat img = cv::imread("input.jpg");
cv::cvtColor(img, img, cv::COLOR_BGR2RGB);
auto img_tensor = torch::from_blob(img.data, {1, img.rows, img.cols, 3}, torch::TensorOptions().dtype(torch::kByte)).to(torch::kCUDA);
img_tensor = img_tensor.permute({0,3,1,2});
img_tensor = img_tensor.toType(torch::kFloat);

/* CUDAFloatType { 1, 3, h, w } -> CUDAFloatType { 1, 1, h, w }*/
torch::NoGradGuard no_grad;
auto result = net.forward({img_tensor}).toTensor();

/* CUDAFloatType { 1, 1, h, w } -> CPUByteType { h, w } */
result = result.squeeze().sigmoid().mul(255.0).toType(torch::kByte).to(torch::kCPU);

/* CPUByteType { h, w } -> cv::Mat { h, w, 1 } */
cv::Mat img_C1;
img_C1.create(cv::Size(img.cols, img.rows), CV_8UC1);
memcpy(img_C1.data, result.data_ptr(), result.numel() * sizeof(torch::kByte));
cv::imwrite("out.png", img_C1);

三通道

神经网络输出的 tensor 的通道数为 3。

#include <torch/torch.h>
#include <torch/script.h>

/* 初始化模型 */
torch::jit::script::Module net = torch::jit::load("../models/poolnet.pt");
net.to(torch::kCUDA);
net.eval();

/* cv::Mat { h, w, 3 } -> CUDAFloatType { 1, 3, h, w } */
cv::Mat img = cv::imread("input.jpg");
cv::cvtColor(img, img, cv::COLOR_BGR2RGB);
auto img_tensor = torch::from_blob(img.data, {1, img.rows, img.cols, 3}, torch::TensorOptions().dtype(torch::kByte)).to(torch::kCUDA);
img_tensor = img_tensor.permute({0,3,1,2});
img_tensor = img_tensor.toType(torch::kFloat);

/* CUDAFloatType { 1, 3, h, w } -> CUDAFloatType { 1, 3, h, w }*/
torch::NoGradGuard no_grad;
auto result = net.forward({img_tensor}).toTensor();

/* CUDAFloatType { 1, 3, h, w } -> CPUByteType { h, w, 3 } */
result = result.squeeze().sigmoid().mul(255.0).toType(torch::kByte).permute({1,2,0}).to(torch::kCPU);

/* CPUByteType { h, w, 3 } -> cv::Mat { h, w ,3 } */
cv::Mat img_C3;
img_C3.create(cv::Size(img.cols, img.rows), CV_8UC3);
memcpy(img_C3.data, result.data_ptr(), result.numel() * sizeof(torch::kByte));
cv::imwrite("out.png", img_C3);

CUDA2CPU

把 tensor 从 GPU 上转移到 CPU，操作是

result = result.to(torch::kCPU);

这样做的目的有两点

• GPU 上 tesnor 的处理和 CPU 主程序是异步的，to(torch::kCPU) 可以强制 tensor 先被 GPU 处理完毕然后再被 CPU 主程序访问。
• 主程序跑在 CPU 上，无法访问 GPU 上的 tensor 的 data_ptr()。