C++ 调用pytorch进行inference

最新推荐文章于 2024-07-16 19:19:12 发布

徐梓航

最新推荐文章于 2024-07-16 19:19:12 发布

阅读量965

点赞数

分类专栏： Pytorch学习文章标签： C++ libtorch Opencv

本文链接：https://blog.csdn.net/xuyunyunaixuexi/article/details/103387495

版权

本文介绍了如何在C++中利用libtorch调用预先训练的PyTorch模型进行推理，特别是在分割网络上的应用。通过对比，发现C++的推理速度显著快于PyTorch环境，例如在测试296张图片时，C++仅需33秒，而PyTorch则需要140秒。这展示了C++在模型部署和效率上的优势。

摘要由CSDN通过智能技术生成

int main() 
{
	int img_size = 224;
	torch::DeviceType device_type = at::kCUDA;
	cout << "use cuda" << endl;
	// Deserialize the ScriptModule from a file using torch::jit::load().
	Module module = torch::jit::load("G:/pytorch/model.pt");
	std::cout << "ok\n";
	module.to(device_type);
	assert(module != nullptr);
	std::cout << "ok\n";
	
	Mat src, image, float_image;
	src = imread("G:/pytorch/predict-app/data/1.jpg"); //读图

	resize(src, image, Size(img_size, img_size));


	cvtColor(image, image, CV_BGR2RGB);//bgr->rgb
	image.convertTo(float_image, CV_32F, 1.0 / 255);//归一化到[0,1]区间


	auto img_tensor = torch::from_blob(float_image.data, { 1, 224, 224, 3 });

	img_tensor = img_tensor.permute({ 0, 3, 1, 2 });

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

	i