使用 TensorRT 进行推理的系统讲解

使用 TensorRT 进行推理需要几个主要步骤，从创建 Runtime 到最终的推理。这些步骤包括：

1. 创建 Runtime

2. 反序列化引擎

3. 创建 Execution Context

4. 分配内存

5. 图像前处理

6. 执行推理

7. 结果后处理

以下是对每一步的详细解释：

1. 创建 Runtime

目的

nvinfer1::IRuntime 是 TensorRT 进行推理的核心组件。创建 Runtime 的目的是提供一个上下文环境来管理推理引擎和执行上下文。

实现

nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(gLogger);

注意事项

• 确保 gLogger 被正确初始化，gLogger 是 TensorRT 提供的日志记录器，用于记录和调试信息。

• 需要在程序结束时销毁 runtime 以释放资源。

2. 反序列化引擎

目的

引擎是经过优化的推理模型。反序列化引擎的目的是将二进制格式的引擎文件加载到内存中，以便执行推理。

实现

std::ifstream engine_file(engine_file_path, std::ios::binary); engine_file.seekg(0, std::ifstream::end); size_t fsize = engine_file.tellg(); engine_file.seekg(0, std::ifstream::beg); std::vector<char> engine_data(fsize); engine_file.read(engine_data.data(), fsize); nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(engine_data.data(), fsize, nullptr);

注意事项

• 确保引擎文件存在且可读。

• deserializeCudaEngine 函数将引擎文件内容反序列化为可用的推理引擎对象。

3. 创建 Execution Context

目的

Execution Context 是引擎在推理过程中实际使用的上下文。它负责管理推理过程中所有的操作，包括内存分配和内核执行。

实现

nvinfer1::IExecutionContext* context = engine->createExecutionContext();

注意事项

• 每个 Execution Context 对象只能用于一个线程中。

• 需要在程序结束时销毁 context 以释放资源。

4. 分配内存

目的

为模型的输入和输出分配 GPU 内存，以便在推理过程中进行数据传输和计算。

实现

void* buffers[2]; cudaMalloc(&buffers[0], input_size * sizeof(float)); // 输入 cudaMalloc(&buffers[1], output_size * sizeof(float)); // 输出 cudaStream_t stream; cudaStreamCreate(&stream);

注意事项

• 确保分配的内存大小和输入输出的大小匹配。

• 在程序结束时需要释放分配的 GPU 内存。

5. 图像前处理

目的

将输入图像转换为模型所需的格式，并将其传输到 GPU 内存中。

实现

注意事项

• 确保输入图像大小和模型要求的输入大小一致。

• 使用 CUDA 异步内存复制函数 cudaMemcpyAsync，可以提升性能。

6. 执行推理

目的

使用 Execution Context 执行推理，得到模型的输出结果。

实现

context->enqueueV2(buffers, stream, nullptr);

注意事项

• enqueueV2 是异步执行，需要使用 CUDA 流来管理同步。

7. 结果后处理

目的

将 GPU 内存中的推理结果传回主机内存，并进行后处理以得到最终的输出结果。

实现

注意事项

• 确保内存复制和流同步，以确保结果数据的完整性。

示例代码汇总