使用LibTorch进行C++调用pytorch模型方式

import torch
import torchvision.models as models
from PIL import Image
import numpy as np
from torchvision import transfORMs

model_resnet = models.resnet50(pretrained=True).cuda()

# model_resnet.load_state_dict(torch.load("resnet_Epoch_4_Top1_99.75845336914062.pkl"))
model_resnet.eval()
# 自己选择任意一张图片，并将它的路径写在open方法里，用来读取图像，我这里路径就是‘111.jpg'了
image = Image.open("111.jpg").convert('RGB')

tf = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        # transforms.Normalize(mean=[0.5]*3, std=[0.5]*3)
])

img = tf(image)
img = img.unsqueeze(dim=0)
print(img.shape)
input = torch.rand(1, 3, 224, 224).cuda()

traced_script_module_resnet = torch.jit.trace(model_resnet, input)

output = traced_script_module_resnet(img.cuda())
print(output.shape)
pred = torch.argmax(output, dim=1)
print(pred)
traced_script_module_resnet.save("model_resnet_jit_cuda.pt")

#include <torch/script.h> // One-stop header.
#include <opencv2/opencv.hpp>
#include <iOStream>
#include <memory>

//https://pytorch.org/tutorials/advanced/cpp_export.html

std::string image_path = "../../111.jpg";

int main(int arGC, const char* argv[]) {

	// Deserialize the ScriptModule from a file using torch::jit::load().
	//std::shared_ptr<torch::jit::script::Module> module = torch::jit::load("../../model_resnet_jit.pt");
	using torch::jit::script::Module;
	Module module = torch::jit::load("../../model_resnet_jit_cuda.pt");
	module.to(at::kCUDA);

	//assert(module != nullptr);
	//std::cout << "ok\n";

	//输入图像
	auto image = cv::imread(image_path, cv::ImreadModes::IMREAD_COLOR);
	cv::cvtColor(image, image, cv::COLOR_BGR2RGB);
	cv::Mat image_transfomed;
	cv::resize(image, image_transfomed, cv::Size(224, 224));

	// 转换为Tensor
	torch::Tensor tensor_image = torch::from_blob(image_transfomed.data,
		{ image_transfomed.rows, image_transfomed.cols,3 }, torch::kByte);
	tensor_image = tensor_image.permute({ 2,0,1 });
	tensor_image = tensor_image.toType(torch::kFloat);
	tensor_image = tensor_image.div(255);
	tensor_image = tensor_image.unsqueeze(0);
	tensor_image = tensor_image.to(at::kCUDA);

	// 网络前向计算
	at::Tensor output = module.forward({ tensor_image }).toTensor();
	//std::cout << "output:" << output << std::endl;

	auto prediction = output.argmax(1);
	std::cout << "prediction:" << prediction << std::endl;

	int maxk = 3;
	auto top3 = std::get<1>(output.topk(maxk, 1, true, true));

	std::cout << "top3: " << top3 << '\n';

	std::vector<int> res;
	for (auto i = 0; i < maxk; i++) {
		res.push_back(top3[0][i].item().toInt());
	}
	for (auto i : res) {
		std::cout << i << " ";
	}
	std::cout << "\n";

	system("pause");
}