深度学习之模型部署【4】-libtorch入门 - pytorch部署torchscript 以及 c++ libtorch 调用 pytorch 模型

pytorch 部署 torchscript

from torchvision.models import resnet34
import torch.nn.functional as F
import torch.nn as nn
import torch
import cv2

#读取一张图片，并转换成[1,3,224,224]的float张量并归一化
image = cv2.imread("flower.jpg")
image = cv2.resize(image,(224,224))
input_tensor = torch.tensor(image).permute(2,0,1).unsqueeze(0).float()/225.0

#定义并加载resnet34模型在imagenet预训练的权重
model = resnet34(pretrained=True)
model.eval()
#查看模型预测该付图的结果
output = model(input_tensor)
output = F.softmax(output,1)
print("模型预测结果为第{}类，置信度为{}".format(torch.argmax(output),output.max()))

#生成pt模型，按照官网来即可
model=model.to(torch.device("cpu"))
model.eval()
var=torch.ones((1,3,224,224))
traced_script_module = torch.jit.trace(model, var)
traced_script_module.save("resnet34.pt")

c++ libtorch 调用 pytorch 模型

cpu 版测试成功，gpu 可能是 libtorch gpu 版本不对，后期再研究吧

#include<opencv2/opencv.hpp>
#include <torch/torch.h>
#include <torch/script.h> 

int main_succ1()
{
    //定义使用cuda
    bool s = torch::cuda::is_available;
    std::cout << s << std::endl;
    auto device = torch::Device(torch::kCUDA, 0);
    //读取图片
    auto image = cv::imread("C:/Users/yansh/source/repos/data/2.jpg");
    //缩放至指定大小
    cv::resize(image, image, cv::Size(224, 224));
    //转成张量
    auto input_tensor = torch::from_blob(image.data, { image.rows, image.cols, 3 }, torch::kByte).permute({ 2, 0, 1 }).unsqueeze(0).to(torch::kFloat32) / 225.0;
    //input_tensor.to(device);
    //加载模型
    auto model = torch::jit::load("C:/Users/yansh/source/repos/data/resnet34.pt");
    //model.to(device);
    //model.eval();
    //前向传播
    auto output = model.forward({input_tensor}).toTensor();
    output = torch::softmax(output, 1);
    std::cout << "模型预测结果为第" << torch::argmax(output) << "类，置信度为" << output.max() << std::endl;
    system("pause");
    return 0;
}

测试代码2

int main() {
    //torch::Device device(torch::kCUDA);
    auto device = torch::Device(torch::kCUDA, 0);
    // Deserialize the ScriptModule from a file using torch::jit::load()
    torch::jit::script::Module module = torch::jit::load("C:/Users/yansh/source/repos/data/resnet34.pt");
    //module.to(device);

    // Create a vector of inputs
    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(torch::ones({ 1, 3, 224, 224 }));    //.to(device)

    // Exectute the model
    at::Tensor output = module.forward(inputs).toTensor();

    std::cout << output.slice(/*dims=*/1, /*start=*/0, /*end=*/5) << '\n';

    std::cout << "ok\n";
    system("pause");
}