最近在C++调用pytorch模型的时候遇到了一点小问题,特此记录。
基于python和pytorch训练了一个分类模型,保存为pkl文件;因C++和Python读取的文件方式不同,所以在使用C++调用之间把训练生成的模型文件转换成Torch Script,即生成一个pt文件,这样就可以使用C++进行调用。
生成的pkl和pt文件在python下测试,结果保持一致,证明模型没有问题,很好,继续。
在C++下调用pt文件,发现测试结果与python版本不一致,证明喂给模型的数据有问题。
以下为python和C++的代码片段:
python:
from torchvision import transformsimport os
from PIL import Image
import torch
import torch.nn.functional as F
img_path=D:\\'
data_transforms=transforms.Compose([transforms.Resize(128),transforms.ToTensor() ]) #resize,并将图片转换成tensor
for file in os.listdir(img_path):
imagepath = img_path+"\\"+file image = Image.open(imagepath) #load图片
imgblob = data_transforms(image).unsqueeze(0) #增加维度
torch.no_grad()
pre = F.softmax(model(imgblob.cuda()))
prediction = torch.argmax(pre).item()print(prediction)
C++:
cv::Mat temImage = cv::imread("D:\\1.jpg");
cv::resize(temImage,temImage, cv::Size(128, 128));
torch::Tensor tensor_image = torch::from_blob(temImage.data, { 1,temImage.rows, temImage.cols,3 }, torch::kByte);
tensor_image = tensor_image.permute({ 0,3,1,2 });
tensor_image = tensor_image.toType(torch::kFloat);
auto tensor_image_Tmp = torch::autograd::make_variable(tensor_image, false); //不需要梯度
tensor_image = tensor_image_Tmp.to(torch::kCUDA);
torch::Tensor result = module->forward({ tensor_image }).toTensor();
auto max_result = result.max(1, true);
auto max_index = std::get<1>(max_result).item();
Python接口:图片resize到128X128,然后将图片转换成tensor.
C++接口:图片resize到128X128,然后将图片转换成tensor等.
问题出在:
1:python版,ptorch是用pil读取的图像,C++是用opencv读取的图像,二者读取的图片通道顺序不一样,pil是RGB,opencv是 BGR,所以为了保持一致,需要在resize后加上
cvtColor(temImage, temImage, CV_BGR2RGB); // bgr -> rgb
2:transforms.ToTensor() 将 PIL.Image/numpy.ndarray 数据进转化为torch.FloadTensor,并归一化到[0, 1.0],C++这边只是缺少归一化这个步骤:
temImage.convertTo(temImage, CV_32F, 1.0 / 255); //归一化到[0,1]区间
修改后的代码:
cv::Mat temImage = cv::imread("D:\\1.jpg");
resize(temImage, temImage, cv::Size(128, 128)); // resize 图像
cvtColor(temImage, temImage, CV_BGR2RGB); // bgr -> rgb
temImage.convertTo(temImage, CV_32F, 1.0 / 255); //归一化到[0,1]
auto img_tensor = torch::CPU(torch::kFloat32).tensorFromBlob(temImage.data, {1, temImage.rows, temImage.cols, 3});
img_tensor = img_tensor.permute({0, 3, 1, 2}); //调换顺序变为torch输入的格式 1,3,128,128
auto img_var = torch::autograd::make_variable(img_tensor, false); //不需要梯度
torch::Tensor result = module->forward({ img_var.to(at::kCUDA) }).toTensor(); //前向传播
auto max_result = result.max(1, true);
auto max_index = std::get<1>(max_result).item();
另:pil和opencv在图像处理方面有很多不一样,比如画一条直线,二者用到的插值方法就不同,当处理的图片较大时影响不大,但是图片很小时,就要格外注意。