libtorch学习第一

一、通过python保存模型

from torchvision.models import resnet34
from torchvision.models import ResNet34_Weights

import torch.nn.functional as F
import torch.nn as nn
import torch
import cv2

if __name__ == '__main__':
    img = cv2.imread("../lena.jpg")
    img = cv2.resize(img, (224, 224))
    inputTensor = torch.tensor(img).permute(2,0,1).unsqueeze(0).float()/255.0

    model = resnet34(weights=ResNet34_Weights.DEFAULT)
    model.eval()

    output = model(inputTensor)
    output = F.softmax(output, 1)
    print(f'reuslut class: {torch.argmax(output)}, score: {output.max()}')

    model = model.to(torch.device("cpu"))
    model.eval()
    var = torch.ones((1,3,244,244))
    traced_script_module = torch.jit.trace(model, var)
    traced_script_module.save("resnet34.pt")
    

结果

二、通过c++推理

假设环境已经配好

#include
#include
#include
#include

int main()
{
	/*torch::Tensor output = torch::randn({ 32,2 });
	std::cout << output << std::endl;*/

	cv::Mat img = cv::imread("F:/docker_python/lena.jpg");
	std::cout << img.size << std::endl;

	cv::resize(img, img, cv::Size(224, 224));

	auto inputTensor = torch::from_blob(img.data, { img.rows, img.cols, 3 }, torch::kByte).permute({ 2,0,1 }).unsqueeze(0).to(torch::kFloat32)/ 255.0;

	//inputTensor.to(torch::Device(torch::kCUDA, 0));

	auto model = torch::jit::load("F:\\docker_python\\torch_study\\resnet34.pt");
	
	//model.to(torch::Device(torch::kCUDA, 0));

	//
	model.eval();
	auto output = model.forward({ inputTensor }).toTensor();
	std::cout << "output.shape: ";
	for (auto size : output.sizes()) std::cout << size << " ";
	std::cout<<std::endl; // [1,1000]

	output = torch::softmax(output, 1);
	std::cout << "output.shape: ";
	for (auto size : output.sizes()) std::cout << size << " ";
	std::cout << std::endl; // [1,1000]

	std::cout << "result class: " << torch::argmax(output) << ", score: " << output.max() << std::endl;

	return 0;
}

结果