用TensorRT加速你的YOLOv5：Windows C++推理部署实战（附完整项目配置）

用TensorRT加速YOLOv5：Windows C++推理部署全流程解析

在计算机视觉领域，YOLOv5因其出色的实时检测性能广受欢迎。但当我们需要将训练好的模型部署到实际生产环境时，Python的解释执行往往难以满足性能要求。这时，TensorRT作为NVIDIA推出的高性能推理引擎，能够显著提升模型执行效率。本文将带你从零开始，将一个PyTorch训练的YOLOv5模型转换为TensorRT引擎，并集成到C++应用程序中。

1. 环境准备与模型转换

在开始之前，我们需要确保开发环境配置正确。以下是必需的组件：

• Windows 10/11 64位系统
• NVIDIA显卡（支持CUDA）
• Visual Studio 2019或更高版本
• CUDA 11.x和对应版本的cuDNN
• TensorRT 8.x

提示：务必保持CUDA、cuDNN和TensorRT版本匹配，这是后续步骤成功的关键。

首先，我们需要将训练好的YOLOv5 PyTorch模型(.pt)转换为ONNX格式：

import torch from models.experimental import attempt_load # 加载训练好的模型 model = attempt_load('yolov5s.pt', map_location='cpu') # 设置输入张量尺寸 input_tensor = torch.randn(1, 3, 640, 640) # 导出为ONNX torch.onnx.export( model, input_tensor, 'yolov5s.onnx', opset_version=12, input_names=['images'], output_names=['output'], dynamic_axes={ 'images': {0: 'batch'}, 'output': {0: 'batch'} } )

转换过程中常见的问题及解决方案：

2. TensorRT引擎构建

获得ONNX模型后，我们需要使用TensorRT的builder工具将其转换为优化的推理引擎。这里介绍两种方法：使用trtexec命令行工具和编程方式构建。

2.1 使用trtexec快速转换

trtexec是TensorRT自带的实用工具，适合快速原型开发：

trtexec --onnx=yolov5s.onnx --saveEngine=yolov5s.engine --fp16 --workspace=2048

关键参数说明：

• --fp16: 启用FP16精度，可显著提升性能
• --workspace: 设置最大工作空间大小(MB)
• --minShapes/--optShapes/--maxShapes: 定义动态尺寸范围

2.2 编程方式构建引擎

对于需要更多控制的情况，可以使用TensorRT C++ API：

#include <NvInfer.h> #include <NvOnnxParser.h> nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger); const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); nvinfer1::INetworkDefinition* network = builder->createNetworkV2(explicitBatch); nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, logger); parser->parseFromFile("yolov5s.onnx", nvinfer1::ILogger::Severity::kWARNING); nvinfer1::IBuilderConfig* config = builder->createBuilderConfig(); config->setMaxWorkspaceSize(1 << 30); if (builder->platformHasFastFp16()) { config->setFlag(nvinfer1::BuilderFlag::kFP16); } nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);

构建引擎时的优化技巧：

• 层融合：TensorRT会自动融合Conv+BN+ReLU等常见组合
• 精度校准：对于INT8量化，需要提供校准数据集
• 动态形状：合理设置优化配置文件和内存限制

3. C++推理代码实现

有了TensorRT引擎后，我们需要编写C++代码来加载并执行推理。以下是核心代码结构：

3.1 引擎加载与上下文创建

std::ifstream engineFile("yolov5s.engine", std::ios::binary); engineFile.seekg(0, std::ios::end); size_t engineSize = engineFile.tellg(); engineFile.seekg(0, std::ios::beg); std::vector<char> engineData(engineSize); engineFile.read(engineData.data(), engineSize); nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger); nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(engineData.data(), engineSize); nvinfer1::IExecutionContext* context = engine->createExecutionContext();

3.2 内存分配与数据传输

// 获取输入输出绑定信息 int nbBindings = engine->getNbBindings(); std::vector<void*> buffers(nbBindings); for (int i = 0; i < nbBindings; ++i) { nvinfer1::Dims dims = engine->getBindingDimensions(i); size_t size = std::accumulate(dims.d, dims.d + dims.nbDims, 1, std::multiplies<size_t>()); cudaMalloc(&buffers[i], size * sizeof(float)); } // 将输入数据从主机拷贝到设备 cudaMemcpy(buffers[inputIndex], inputData.data(), inputSize * sizeof(float), cudaMemcpyHostToDevice);

3.3 执行推理与结果处理

context->executeV2(buffers.data()); // 将输出数据从设备拷贝回主机 std::vector<float> outputData(outputSize); cudaMemcpy(outputData.data(), buffers[outputIndex], outputSize * sizeof(float), cudaMemcpyDeviceToHost); // 解析YOLOv5输出 std::vector<Detection> detections; parseYOLOv5Output(outputData, detections);

4. 前后处理优化

在实际应用中，前后处理往往成为性能瓶颈。以下是几种优化策略：

4.1 图像预处理加速

传统CPU预处理：

cv::Mat image = cv::imread("input.jpg"); cv::resize(image, image, cv::Size(640, 640)); image.convertTo(image, CV_32F, 1.0/255.0);

优化后的GPU预处理：

void preprocessGPU(const cv::Mat& h_image, float* d_input, cudaStream_t stream) { // 分配设备内存 uchar* d_uchar; cudaMalloc(&d_uchar, h_image.rows * h_image.cols * 3); // 拷贝并转换 cudaMemcpyAsync(d_uchar, h_image.data, h_image.rows * h_image.cols * 3, cudaMemcpyHostToDevice, stream); // 调用CUDA核函数进行归一化和通道重排 preprocessKernel<<<grid, block, 0, stream>>>(d_uchar, d_input, h_image.cols, h_image.rows); }

4.2 后处理优化

YOLOv5的后处理主要包括：

1. 解码边界框坐标
2. 应用置信度阈值
3. 执行非极大值抑制(NMS)

优化后的NMS实现：

__global__ void nmsKernel(Detection* detections, int num_detections, float iou_threshold, int* keep_indices) { // 共享内存存储检测框信息 extern __shared__ float shared_boxes[]; // 每个线程处理一个检测框 int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= num_detections) return; // 加载检测框到共享内存 if (threadIdx.x == 0) { for (int j = 0; j < num_detections; ++j) { shared_boxes[j*5 + 0] = detections[j].x1; // 加载其他坐标... } } __syncthreads(); // 计算IoU并执行抑制 // ... }

5. 性能对比与调优

完成部署后，我们需要评估TensorRT带来的性能提升。以下是典型测试结果：

性能调优的关键点：

1. 批处理大小：适当增大批处理可提高吞吐量，但会增加延迟
2. 精度选择：FP16通常能在精度损失很小的情况下显著提升性能
3. CUDA流：使用多个CUDA流实现流水线并行
4. 内存复用：避免频繁分配释放内存

// 使用CUDA流实现异步执行 cudaStream_t stream; cudaStreamCreate(&stream); while (true) { // 异步预处理 preprocessGPU(image, d_input, stream); // 异步推理 context->enqueueV2(buffers.data(), stream, nullptr); // 异步后处理 postprocessGPU(d_output, detections, stream); cudaStreamSynchronize(stream); }

在实际项目中，我们还需要考虑工程化方面的优化：

• 异常处理：健壮的错误检查和恢复机制
• 日志系统：详细的性能监控和调试信息
• 资源管理：使用RAII模式管理CUDA资源
• 多线程支持：线程安全的TensorRT上下文管理

通过以上步骤，我们成功将YOLOv5模型部署到了Windows C++环境中，并利用TensorRT实现了显著的性能提升。这种部署方式特别适合需要低延迟、高吞吐量的生产环境，如视频分析、工业检测等应用场景。