Libtorch加载训练好的U-Net网络模型获得预测结果

本人使用的环境是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;
}

以下是运行效果：

 千里之行，始于足下。