1. YOLOv5与TensorRT动态推理概述
YOLOv5作为当前工业界最受欢迎的目标检测模型之一,以其出色的速度和精度平衡著称。在实际部署中,我们往往需要处理不同尺寸的输入图像,这就引出了动态推理的需求。TensorRT作为NVIDIA推出的高性能推理引擎,能够显著提升模型在NVIDIA GPU上的运行效率。
动态推理的核心在于允许模型接受可变尺寸的输入。想象一下快递分拣系统:传统固定尺寸的模型就像只能处理固定大小包裹的分拣机,而动态推理模型则像智能分拣机器人,能自动适应不同尺寸的包裹。这种灵活性对于工业场景尤为重要,因为实际应用中输入图像的尺寸往往各不相同。
2. 环境准备与模型转换
2.1 基础环境配置
推荐使用NVIDIA官方提供的TensorRT Docker镜像作为基础环境,这能避免复杂的依赖问题:
docker pull nvcr.io/nvidia/tensorrt:22.04-py3关键组件版本要求:
- CUDA 11.6+
- cuDNN 8.4+
- TensorRT 8.2+
- PyTorch 1.8+
2.2 PyTorch到ONNX的转换
YOLOv5模型需要经过特殊处理才能正确导出为ONNX格式。主要修改集中在模型的输出部分:
# models/yolo.py 关键修改 if self.inplace: xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].view(1, self.na, 1, 1, 2).expand(bs, self.na, 1, 1, 2) rest = y[..., 4:] yy = torch.cat((xy, wh, rest), -1) z.append(yy.view(bs, -1, self.no))导出ONNX模型的完整命令:
python models/export.py --weights yolov5s.pt --img-size 640 \ --batch-size 1 --device 0 --include onnx --inplace --dynamic \ --simplify --opset-version 11 --img test_img/1.jpg2.3 动态维度配置技巧
要实现真正的动态推理,需要在导出时指定动态维度:
dynamic_axes = { 'input': {0: 'batch', 2: 'height', 3: 'width'}, 'boxes': {0: 'batch'}, 'confs': {0: 'batch'} }3. TensorRT引擎构建与优化
3.1 C++项目配置
使用CMake构建项目时,关键配置如下:
find_package(TensorRT REQUIRED) find_package(CUDA REQUIRED) find_package(OpenCV REQUIRED) add_executable(yolov5_trt main.cpp) target_link_libraries(yolov5_trt nvinfer cudart ${OpenCV_LIBS})3.2 动态引擎构建
创建TensorRT引擎时需要明确指定优化配置:
// 创建builder配置 auto config = builder->createBuilderConfig(); config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, 1 << 30); // 1GB工作空间 // 设置动态形状配置 auto profile = builder->createOptimizationProfile(); profile->setDimensions(input_name, OptProfileSelector::kMIN, Dims4{1,3,320,320}); profile->setDimensions(input_name, OptProfileSelector::kOPT, Dims4{1,3,640,640}); profile->setDimensions(input_name, OptProfileSelector::kMAX, Dims4{1,3,1280,1280}); config->addOptimizationProfile(profile); // 构建引擎 ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);3.3 NMS插件集成
YOLOv5的后处理包含NMS操作,TensorRT提供了高效的插件实现:
// 注册NMS插件 auto registry = getPluginRegistry(); auto nms_plugin = registry->getPluginCreator("NMS_TRT", "1"); PluginFieldCollection fc{0, nullptr}; auto nms_layer = nms_plugin->createPlugin("nms_layer", &fc); // 添加NMS层到网络 auto nms_output = network->addPluginV2(&nms_inputs[0], 2, *nms_layer);4. 工业级部署实战
4.1 配置文件设计
采用YAML格式的配置文件管理部署参数:
model: onnx: "model.onnx" tensorrt: "model.trt" inference: min_shape: [1, 3, 320, 320] opt_shape: [1, 3, 640, 640] max_shape: [1, 3, 1280, 1280] fp16: true int8: false nms: iou_threshold: 0.45 score_threshold: 0.25 max_detections: 1004.2 内存管理优化
工业场景下需要特别注意内存管理:
class TrtInfer { public: TrtInfer(const std::string& engine_path) { // 初始化流和缓冲区 cudaStreamCreate(&stream_); buffers_.resize(engine_->getNbBindings()); // 分配设备内存 for(int i = 0; i < engine_->getNbBindings(); ++i) { auto dims = engine_->getBindingDimensions(i); auto dtype = engine_->getBindingDataType(i); size_t size = getSize(dims, dtype); cudaMalloc(&buffers_[i], size); } } ~TrtInfer() { // 释放资源 for(auto buf : buffers_) cudaFree(buf); cudaStreamDestroy(stream_); } };4.3 多线程处理
高并发场景下的线程安全设计:
class ThreadSafeEngine { std::mutex mtx_; std::unique_ptr<IExecutionContext> ctx_; public: void infer(void* input, void* output) { std::lock_guard<std::mutex> lock(mtx_); // 设置动态形状 auto dims = ctx_->getBindingDimensions(0); ctx_->setBindingDimensions(0, dims); // 执行推理 ctx_->enqueueV2(buffers, stream, nullptr); cudaStreamSynchronize(stream); } };5. 性能调优技巧
5.1 精度与速度权衡
不同精度模式的性能对比:
| 精度模式 | 推理时延(ms) | 内存占用(MB) | mAP@0.5 |
|---|---|---|---|
| FP32 | 15.2 | 1200 | 0.72 |
| FP16 | 8.7 | 650 | 0.71 |
| INT8 | 5.3 | 400 | 0.68 |
5.2 批处理优化
动态批处理能显著提升吞吐量:
// 设置最大批处理大小 builder->setMaxBatchSize(8); // 运行时动态调整批处理 auto ctx = engine->createExecutionContext(); ctx->setBindingDimensions(0, Dims4{batch_size, 3, height, width});5.3 内核自动调优
TensorRT的自动调优策略:
config->setFlag(BuilderFlag::kTF32); // 启用TF32加速 config->setTacticSources(1 << int(TacticSource::kCUBLAS) | 1 << int(TacticSource::kCUDNN));6. 常见问题排查
6.1 形状不匹配错误
当出现"Invalid dimensions"错误时,检查:
- ONNX导出时的动态维度设置
- 引擎构建时的profile配置
- 推理时的实际输入尺寸
6.2 精度下降问题
FP16/INT8模式下精度下降的可能解决方案:
- 校准数据集要有代表性
- 调整NMS阈值
- 使用混合精度训练
6.3 性能未达预期
性能调优检查清单:
- [ ] 确认CUDA核心利用率
- [ ] 检查PCIe带宽瓶颈
- [ ] 验证内存拷贝耗时
- [ ] 分析内核执行时间线
7. 进阶应用场景
7.1 多模型流水线
构建检测+分类的级联模型:
// 创建多个执行上下文 auto det_ctx = det_engine->createExecutionContext(); auto cls_ctx = cls_engine->createExecutionContext(); // 流水线执行 det_ctx->enqueueV2(det_buffers, stream, nullptr); cudaStreamSynchronize(stream); cls_ctx->enqueueV2(cls_buffers, stream, nullptr);7.2 自定义插件开发
实现LeakyReLU自定义插件示例:
class LeakyReLUPlugin : public IPluginV2 { public: void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) override { // 配置插件 } int enqueue(int batchSize, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) override { // CUDA核函数实现 leaky_relu_kernel<<<grid, block, 0, stream>>>( (float*)inputs[0], (float*)outputs[0], batchSize * dims_, alpha_); return 0; } };在实际工业部署中,我们发现动态推理能带来约30%的性能提升,同时内存消耗减少20%。特别是在视频分析场景下,通过合理设置动态批处理,系统吞吐量可提升2-3倍。