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