libtorch 调用pytorch模型

我测试的速度没有提升


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

	//size_t len = url.length();//获取字符串长度


	char sBuf[1024];
	char *ptr;
	if (GetModuleFileNameA(NULL, sBuf, sizeof(sBuf)))
	{
		ptr = strrchr(sBuf, '\\');
		if (ptr)
			*ptr = '\0';
		SetCurrentDirectoryA(sBuf);
	}


	torch::DeviceType device_type;

	if (torch::cuda::is_available()) {
		device_type = torch::kCUDA;
	}
	else {
		device_type = torch::kCPU;
	}
	torch::Device device(device_type);

try
	{
		std::shared_ptr module = torch::jit::load("../models/yolo_model.pt");

		assert(module != nullptr);

		//module->to(at::kCUDA);
		module->to(device);
		torch::Tensor img_tensor = torch::rand({ 1, 3,352,352 }).to(device);

		//img_tensor.to(at::kCUDA)
		//auto img_var = torch::autograd::make_variable(img_tensor, false).to(device);
		for (int i = 0; i < 20;i++) {
			DWORD start, end;

			start = GetTickCount();
			auto output = module->forward({ img_tensor });
			//at::Tensor output = module->forward({ img_tensor }).toTensor();
			end = GetTickCount() - start;
			cout << "time " << end << "ms!" << endl;
		
		}
		std::cout << "load ok---------" << endl;
	}
	catch (string &e)
	{
		std::cout << e<< endl;
	}
	
	return 0;

 

你可能感兴趣的:(c++,torch)