博客

1. 加载pointNet引擎文件并创建引擎对象

// 加载引擎文件
	std::string engine_name = "./pointnet.engine";
	std::ifstream file(engine_name, std::ios::binary);
	if (!file.good())
		std::cerr << "文件无法打开，请确定文件是否可用！" << std::endl;
	size_t size = 0;
	file.seekg(0, file.end);	// 将读指针从文件末尾开始移动0个字节
	size = file.tellg();	// 返回读指针的位置，此时读指针的位置就是文件的字节数
	file.seekg(0, file.beg);	// 将读指针从文件开头开始移动0个字节
	char *modelStream = new char[size];
	file.read(modelStream, size);
	file.close();
	// 创建引擎对象
	Logger logger;
	IRuntime* runtime = createInferRuntime(logger);
	ICudaEngine* engine = runtime->deserializeCudaEngine(modelStream, size);
	IExecutionContext* context = engine->createExecutionContext();
	delete[] modelStream;

2. 申请存储空间

// 在主机和GPU上申请输入输出的存储空间
	const int inputNum = B * N * C;
	const int outputNum = B * K;
	float *data = (float*)malloc(inputNum * sizeof(float));
	float *prob = (float*)malloc(outputNum * sizeof(float));
	float* buffers[2];
	const int inputIndex = engine->getBindingIndex(inputName);
	const int outputIndex = engine->getBindingIndex(outputName);
	CUDA_CHECK(cudaMalloc((void**)&buffers[inputIndex], inputNum * sizeof(float)));
	CUDA_CHECK(cudaMalloc((void**)&buffers[outputIndex], outputNum * sizeof(float)));

3. 加载点云文件并进行推理

// 获取所有文件名
	std::string pc_dir = "./pointCloud";
	std::vector<std::string> file_names;
	read_files_in_dir(pc_dir.c_str(), file_names);
	// 创建cuda流
	cudaStream_t stream;
	CUDA_CHECK(cudaStreamCreate(&stream));
	// 遍历所有文件，一次取batchsize个数在显存中推理后返回结果
	int batchNum = 0;
	for (int f = 0; f < (int)file_names.size(); f += B) {
		// 读取一个批次的数据
		readPC(pc_dir, file_names, batchNum, B, C, N, data);	

		// 传输数据并推理
		CUDA_CHECK(cudaMemcpyAsync(buffers[inputIndex], data, inputNum * sizeof(float), cudaMemcpyHostToDevice, stream));	// 异步传输数据到GPU
		context->enqueue(B, (void**)buffers, stream, nullptr);	// 使用异步推理
		CUDA_CHECK(cudaMemcpyAsync(prob, buffers[1], outputNum * sizeof(float), cudaMemcpyDeviceToHost, stream));	// 传回结果数据
		CUDA_CHECK(cudaStreamSynchronize(stream));	// 同步该流，等待该流上的命令都完成

		// 做预测
		for (int i = 0; i < B; ++i) {
			int pos = batchNum * B + i;
			if (pos < file_names.size())
				std::cout << file_names[pos] << ":" << types[find_max(prob + K * i, K)] << std::endl;
		}

		++batchNum;
	}

4. 释放资源和空间

//  释放资源和空间
 	 >destroy();
  	runtime->destroy();
  	CUDA_CHECK(cudaFree(buffers[inputIndex]));
   	CUDA_CHECK(cudaFree(buffers[outputIndex]));
	free(data);
 	free(prob);