windows 10 c++(libtorch)调用pytorch训练好的分割模型

c++调用训练好的分割模型主要步骤如下：

1、将训练好的.pth 模型转换为.pt模型
2、opencv读取要推理的图像
3、model前向传播，获得推理结果
4、将推理结果转为opencv，查看推理结果。

开发环境：

windows10
vs 2015
pytorch 1.2.0 cpu
opencv 4.3.0
libtorch 1.3.0 cpu

配置开发环境时遇到的坑：
1、vs2015不支持最新的libtorch
2、ptorch与liborch 使用的cpu或者gpu版本要统一，否则出现如下错误：

将训练好的.pth模型模型转为.pt模型

见上一篇文章：
pytorch训练好的.pth模型转换为.pt

2.2、c++代码如下：

#include "torch/script.h" // One-stop header. 
#include  
#include  
#include  
using namespace cv; 
using namespace std; 
int main() 
{  
 Deserialize the ScriptModule from a file 
 torch::jit::script::Module module = torch::jit::load("caoxie_weight.pt");  
 //assert(module != nullptr);  
 auto image = imread("D:\\UNET\\data\\org\\2.bmp"); 
 if (!image.data)  
 {   
 cout << "image imread failed" << endl; 
  }  
  cvtColor(image, image, CV_BGR2RGB);  
  Mat img_transfomed;  
  resize(image, img_transfomed,Size(480, 320));  
  /*cout << img_transfomed.data;*/  
  //img_transfomed.convertTo(img_transfomed, CV_16FC3, 1.0f / 255.0f);  
  //Mat to tensor,   
  torch::Tensor tensor_image = torch::from_blob(img_transfomed.data, {img_transfomed.rows, 	  img_transfomed.cols, img_transfomed.channels()}, torch::kByte);
  tensor_image = tensor_image.permute({2, 0, 1});
  tensor_image = tensor_image.toType(torch::kFloat);
  tensor_image = tensor_image.div(255);
  tensor_image = tensor_image.unsqueeze(0);
 std::vector<torch::jit::IValue> inputs; 
 inputs.push_back(tensor_image);

 torch::Tensor output = module.forward(inputs).toTensor();  
 torch::Tensor output_max = output.argmax(1);
 //tensor to Mat  
 output_max = output_max.squeeze();  
 output_max = output_max.mul(255).to(torch::kU8);  
 output_max = output_max.to(torch::kCPU);
 Mat result_img(Size(480, 320), CV_8UC1);  
 memcpy((void*)result_img.data, output_max.data_ptr(), sizeof(torch::kU8)*output_max.numel());
 imshow("result", result_img);  
 imwrite("result.bmp", result_img);  
 system("pause");
 }

2.3 cmake 编译

cmakelist 如下：

cmake_minimum_required(VERSION 3.0 FATAL_ERROR) 
project(example)
find_package(Torch REQUIRED)
add_executable(example example.cpp)
target_link_libraries(example ${TORCH_LIBRARIES})
set_property(TARGET example PROPERTY CXX_STANDARD 11)

命令行cd到当前目录下，传入libtorch路径：

cmake -DCMAKE_PREFIX_PATH=libtorh路径.. -DCMAKE_BUILD_TYPE=Release -G"Visual Studio 14 Win64" //release版

输出结果：

参考资料：

pytorch官网教程：