当前位置：首页 > news >正文

机器学习笔记 - 在QT/PyTorch/C++ 中加载 TORCHSCRIPT 模型

news 来源：原创 2024/5/8 0:15:46

一、配置QT+libtorch环境

1、首先下载pytorch，注意版本要和你的python下的pytorch版本，否则运行脚本时会报错，我这里用的是windows下的cuda11.6，release版本，有debug和release两个版本可以下载。

2、将压缩包内的include和lib文件夹放到指定位置，我这里是放到了项目文件夹内。

3、因为QT项目用的是CMakeLists.txt文件，所以添加如下代码。

//两个包含路径
include_directories(${CMAKE_SOURCE_DIR}/lib/torch/include)
include_directories(${CMAKE_SOURCE_DIR}/lib/torch/include/torch/csrc/api/include)

//这里可以把所有的包都引入，我比较懒，暂时这里只引了几个
target_link_libraries(QtFFmpegApp2 PRIVATE ${CMAKE_SOURCE_DIR}/lib/torch/lib/asmjit.lib
                                           ${CMAKE_SOURCE_DIR}/lib/torch/lib/c10.lib
                                           ${CMAKE_SOURCE_DIR}/lib/torch/lib/torch.lib
                                           ${CMAKE_SOURCE_DIR}/lib/torch/lib/torch_cpu.lib
                                           ${CMAKE_SOURCE_DIR}/lib/torch/lib/torch_cuda.lib
                                       )

4、测试环境，这里需要注意引入包含文件的时候按照如下写法，应该是slots有冲突。

//torch
#undef slots
#include <torch/torch.h>
#define slots Q_SLOTS


torch::Tensor tensor = torch::rand({ 5,3 });
std::cout << tensor << std::endl;

输出如下

 0.7113  0.9045  0.9027
 0.2154  0.1694  0.4926
 0.1489  0.6318  0.7042
 0.2531  0.5301  0.2073
 0.6361  0.4789  0.3916
[ CPUFloatType{5,3} ]

二、将PyTorch模型转换为Torch脚本

这里直接实用官方的alexnet模型，测试过一些自定义的模型，但是总是失败，还没研究明白，后面再试。

import torch
import torchvision


alexnet = torchvision.models.alexnet(pretrained=True)
dummy_input = torch.rand((1, 3, 224, 224))

traced_alexnet = torch.jit.trace(alexnet, dummy_input)
torch.jit.save(traced_alexnet, "alexnet.pt")

三、在 C++ 中加载脚本模块

#undef slots
#include <torch/torch.h>
#include <torch/script.h>
#define slots Q_SLOTS


void MainWindow::on_pushButton_2_clicked()
{
    //实用opencv读取图片
    cv::Mat img = cv::imread("0.jpg");
    //进行图片缩放
    cv::resize(img, img, cv::Size(224, 224), 0, 0, 1);
    //转为tensor_image 
    auto tensor_image = torch::from_blob(img.data, { img.rows, img.cols, img.channels() }, at::kByte);
    //将tensor的维度换位
    tensor_image = tensor_image.permute({ 2,0,1 });


    auto mean1 = torch::tensor({0.485}).repeat({224, 224});
    auto mean2 = torch::tensor({0.456}).repeat({224, 224});
    auto mean3 = torch::tensor({0.406}).repeat({224, 224});
    auto mean = torch::stack({mean1, mean2, mean3});

    auto std1 = torch::tensor({0.229}).repeat({224, 224});
    auto std2 = torch::tensor({0.224}).repeat({224, 224});
    auto std3 = torch::tensor({0.225}).repeat({224, 224});
    auto std = torch::stack({std1, std2, std3});

    tensor_image = tensor_image / 255.0;
    tensor_image = (tensor_image - mean) / std;

    tensor_image.unsqueeze_(0);

    auto image = tensor_image.to(torch::kFloat);
    
    //如果load失败大概率是因为用Python生成.pt文件时所采用的pytorch的版本和C++采用的Libtorch版本不一致
    torch::jit::script::Module alexnet = torch::jit::load("alexnet.pt");

    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(image);

    auto output = alexnet.forward(inputs).toTensor();
    auto prediction_index = output.argmax();
    
    //这里应该是使用对应的imagenet得到分类名
    //std::cout << "Prediction: " << id2class[prediction_index.item<int>()] << std::endl;
    
    //不过懒得弄，直接打印了索引出来
    std::cout << "Prediction: " << prediction_index.item<int>() << std::endl;
}

这是一只可爱的猫咪，推理完成之后得到的索引是281。在imagenet的分类索引里面是281 n02123045 猫, tabby, tabby cat。分类正确