C++ 调用pytorch进行inference

原创

于 2019-12-04 16:19:55 发布 · 1k 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#C++ #libtorch #Opencv

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

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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));

	inputs.emplace_back((img_tensor).to(at::kCUDA));


	// Execute the mo