C++(Libtorch)加载python(pytorch)网络进行预测（版本、cpu、cuda、runtime error问题）

写在前面：

（1）如果C++中不使用GPU，需保证libtorch的版本与pytorch的版本一致，否则可能会出错，例如出现c++中load网络失败的情况。
（2）如果C++中使用cuda，libtorch版本与pytorch版本可能会不一致，未做统计。
（3）VS为2017版本，系统win10

1. C++中不使用CUDA

我的版本是：
Pytorch：torch-1.6.0-cp36-cp36m-win_amd64.whl（cu102） + torchvision-0.7.0-cp36-cp36m-win_amd64.whl（cu102）
libtorch：libtorch-win-shared-with-deps-1.6.0（cu102）
基本步骤如下：

1.1 Python环境下网络训练保存

首先在python环境下训练网络，然后保存，例如：
torch.save(net.state_dict(), ‘best_model.pth’)

1.2 python环境下生成pt文件

使用trace方法。
（1）新建网络
device = torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’)
device = torch.device(‘cpu’) #不使用GPU
net = UNet(n_channels=1, n_classes=1)
net.to(device=device)
（2）load训练好的权重
dic_path = ‘best_model_202111111946.pth’
net.load_state_dict(torch.load(dic_path))
net.eval()
（3）导入一张图片用以预测。需要将图片转为tensor格式
img = cv2.imread(‘data/testAl/image/9.tif’,-1)
img = img.reshape(1, 1, img.shape[0], img.shape[1])
img_tensor = torch.from_numpy(img)
img_tensor = img_tensor.to(device=device, dtype=torch.float32)
这里，reshape成什么样的维度，取决于网络训练时使用的参数。不导入图片，使用产生随机数来预测也可以，例如：
img = torch.ones(1, 1, 32,32)
img_tensor = img.to(device=device, dtype=torch.float32)
（4）trace并保存pt文件
traced_script_module_resnet = torch.jit.trace(net, img_tensor)
output = traced_script_module_resnet(img_tensor)
traced_script_module_resnet.save(“model_20211112_jit.pt”)
这里，output是网络的输出，可与C++环境预测的输出做对比。

1.3 C++环境下导入pt文件并预测

（1）导入pt文件
torch::jit::script::Module module = torch::jit::load("\model_2111111021_jit.pt");
module.eval();
//module.to(at::kCUDA);
（2）读取图片并转为tensor
Mat imgSrc = imread("…/0.png", -1);
torch::Tensor tensor_image = torch::from_blob(imgSrc.data,{ 1,1,imgSrc.rows, imgSrc.cols }, torch::kByte);
tensor_image = tensor_image.toType(torch::kFloat);
std::vectortorch::jit::IValue inputs;
inputs.push_back(tensor_image);
//tensor_image = tensor_image.to(at::kCUDA);
注意，mat转tensor，需要转成何种维度，取决于网络训练时的输入维度
（3）使用网络预测
at::Tensor outputs = module.forward(inputs).toTensor();
（4）获得输出并转为图片
outputs = outputs[0];
outputs = outputs[0];
cout << outputs.dim() << endl;
Mat outimg(cv::Size(outputs.size(0),outputs.size(1)), CV_32F, outputs.data_ptr());

2. C++中使用CUDA

C++中使用cuda进行网络预测时（在保证vs2017环境部署好的前提下），如果libtorch仍然使用上述版本，则会出现问题。现象如下：
检测cuda和cudnn，通过：
cout << “cuda是否可用：” << torch::cuda::is_available() << endl;
cout << “cudnn是否可用：” << torch::cuda::cudnn_is_available() << endl;

导入pt文件通过：
string modPaht = “202111180132.pt”;
torch::jit::script::Module module = torch::jit::load(modPaht);
module.to(at::kCUDA);
module.eval();
产生数据，并置于cuda，通过，
torch::Tensor tensor = torch::ones({ 1,1,32,32 });
tensor = tensor.to(at::kCUDA);
cout << “data is on:” << tensor.device() << endl;
std::vectortorch::jit::IValue inputs;
inputs.push_back(tensor);

使用网络预测，出现runtime error异常
module.forward(inputs);//出现异常
反复检查后，仍无法解决。后来解决办法是，提高libtorch的版本，将libtorch的版本替换为
libtorch-win-shared-with-deps-1.7.0（cu10.2），问题解决。

3. C++中使用CUDA注意事项

（1）要保证模型和输入数据在同一设备上（CPU或CUDA）
如果使用cuda，在进行module.forward()之前，要将module.to(at::kCUDA);以及输入数据tensor.to(at::kCUDA);
（2）若使用cuda，则复制网络输出Tensor变量至Mat变量前，需先将输出Tensor转移至CPU上
at::Tensor outputs = module.forward({ tensor_image }).toTensor();
//提取预测结果
outputs = outputs[0];
outputs = outputs[0];
outputs = outputs.to(at::kCPU);
Mat outimg(cv::Size(outputs.size(0), outputs.size(1)), CV_32F, outputs.data_ptr());
（3）若python在产生pt文件时，device为‘cpu’，C++中读取pt文件时，在不加限定（torch::jit::load函数中的第二个参数可以限定设备位置）的情况下，会将模型加载到cpu；同样，如果python产生pt文件时，device为‘cuda’设备，在不加限定情况下，会将模型加载到cuda。在做预测时，需保证输入数据的设备位置与模型所在设备位置相同。
（4）如果python产生的是pt文件在cpu上，c++中仍然可以通过module.to(at::kCUDA);将模型加载到cuda上；同理，若pt在cuda上，c++中仍然可以通过module.to(at::kCPU);将模型加载到CPU上。