本人使用的环境是Qt5.14.1+VS2017+libtorch1.11.0-cuda+OpenCV3.4.16,以下介绍各环境用途
Qt用于编程的IDE,个人觉得在Qt中编码,比在VS中要舒服一点。
VS2017 Qt的编译工具。
libtorch pytorch的C++版本,用来项目落地,即使用pytorch进行模型训练,导出的模型使用libtorch进行调用。
配置踩坑看这里:
Windows下PyTorch(LibTorch)配置cuda加速_Terod的博客-CSDN博客_libtorch 加速
【Libtorch】解决“cuda::isavailable= 1”,但是加载GPU模型时仍然报错的问题_watermelongogogo的博客-CSDN博客
OpenCV 进行输入网络前后的数据进行处理,如读取、通道操作、尺寸操作、显示等操作。
这里使用的模型是U-Net网络使用ISIC-2017的数据集训练的模型,需要满足TorchScript的要求,
具体要求参考这里:
三维目标检测:(五)如何将pytorch模型部署到C++工程中及pytorch模型转libtorch模型常见的问题_ZZZZZZZZZZZZZYQ的博客-CSDN博客
以下是具体代码:
#include <torch/script.h>
#include <torch/torch.h>
#include <iostream>
#include <opencv2/opencv.hpp>
using namespace std;
using namespace cv;
int main(int argc, char* argv[])
{
std::cout << "torch::cuda::is_available() = " << torch::cuda::is_available() << std::endl;
torch::Device deviceCUDA(torch::kCUDA, 0);
torch::Device deviceCPU(torch::kCPU);
torch::jit::script::Module module;
try
{
module = torch::jit::load("models\\best_model_jit.pt", deviceCUDA);
}
catch(std::exception exc)
{
std::cout << exc.what() << std::endl;
}
Mat image = imread("images/01.jpg");
if(image.empty())
{
cout <<"image imread failed"<< endl;
}
Mat image_gray;
cvtColor(image, image_gray, COLOR_BGR2GRAY);
Mat image_scalre;
resize(image_gray, image_scalre,Size(512,512));
imshow("src", image_scalre);
Mat t = image_scalre.reshape(1,1);
torch::Tensor tensor = torch::from_blob(t.data, {1,1,512,512}, torch::kU8).toType(torch::kFloat32).to(deviceCUDA);
std::vector<torch::jit::IValue> inputs;
inputs.push_back(tensor);
cout << tensor.itemsize() << endl;
torch::Tensor output = module.forward(inputs).toTensor().to(deviceCPU);
cout << output.itemsize() << endl;
Mat result = Mat(512, 512, CV_8UC1, Scalar(0));
output = output.view({ 1, 1, 512, 512 });
for (int i = 0; i < 512; i++)
{
for (int j = 0; j < 512; j++)
{
if(output[0][0][i][j].item().toFloat() >= 0.5)
{
result.data[512*i + j] = 255;
}
}
}
imshow("res", result);
waitKey(0);
return 0;
}
以下是运行效果:
千里之行,始于足下。