Yolo系列模型的TensorRT-C++推理实践

Yolo系列模型的TensorRT-C++推理实践

在边缘计算设备日益承担复杂视觉任务的今天，如何让YOLO这类目标检测模型跑得更快、更稳、更省资源，已经成为工业落地中的核心命题。尤其是在Jetson Orin、T4服务器等多路视频流并发场景下，Python + PyTorch那“优雅但迟缓”的推理方式很快暴露短板——延迟高、吞吐低、内存占用大，难以满足实时性要求。

于是我们不得不转向更底层、更高性能的技术栈：C++ + TensorRT。这不是简单的语言切换，而是一次从“能用”到“好用”的工程跃迁。NVIDIA的TensorRT作为官方推出的推理优化引擎，凭借层融合、精度校准（FP16/INT8）、内核自动调优等硬核能力，在保持模型精度的同时，常可将YOLOv5或YOLOX的推理耗时压缩至个位数毫秒级别。

本文基于我在Linux环境下部署YOLO系列模型的实际经验，梳理出一套完整的C++集成方案。不讲理论堆砌，只聚焦真实可用的工程实现细节，涵盖环境搭建、引擎加载、预处理设计、异步推理与结果解析全流程，适合希望将算法真正推向生产环境的开发者参考。

快速构建开发环境：使用NVIDIA官方镜像

避免“依赖地狱”的最有效方式，就是直接使用NVIDIA提供的TensorRT容器镜像。它已经集成了CUDA、cuDNN、TensorRT、OpenCV等全套工具链，开箱即用，特别适合在Ampere架构GPU（如3090、A100）或Jetson设备上快速验证。

推荐使用的Docker镜像：

nvcr.io/nvidia/tensorrt:23.09-py3

该版本包含TensorRT 8.6+、CUDA 12.2和cuDNN 8.9，兼容主流硬件。启动容器命令如下：

sudo docker run -it \ --name trt_yolo \ --gpus all \ --shm-size=16g \ --network=host \ -v /path/to/your/code:/workspace \ nvcr.io/nvidia/tensorrt:23.09-py3

关键参数说明：
---gpus all：启用所有GPU资源；
---shm-size：增大共享内存，防止多线程数据传输卡顿；
--v：挂载本地代码目录，便于编辑调试；
---network=host：共享主机网络栈，方便接入RTSP流或HTTP服务。

进入容器后，先验证基础组件是否就绪：

tensorrt --version cmake --version # 建议 >= 3.18

若提示未找到tensorrt命令，可能是PATH问题，可通过dpkg -L tensorrt查找实际路径并添加。

补装OpenCV开发包

虽然镜像自带OpenCV，但通常缺少头文件，导致编译失败。执行以下命令补全：

apt update && apt install libopencv-dev -y

测试是否安装成功：

#include <opencv2/opencv.hpp> int main() { cv::Mat img = cv::Mat::zeros(480, 640, CV_8UC3); return 0; }

能顺利通过g++ test.cpp -o test \pkg-config –cflags –libs opencv4``即表示环境准备完成。

推理流程总览：反序列化 → 上下文创建 → 预处理 → 异步执行

整个C++推理流程遵循TensorRT的标准范式：
加载序列化引擎 → 创建执行上下文 → 图像预处理 → 内存拷贝 → 异步推理 → 后处理输出

我们以一个封装良好的YoloTRT类为主线，逐步拆解其实现逻辑。这种面向对象的设计不仅结构清晰，也便于后续扩展为多实例并发或多模型Pipeline。

自定义日志处理器：别让日志淹没关键信息

所有TensorRT API调用都需要传入一个ILogger实例。默认的日志太啰嗦，我们可以继承nvinfer1::ILogger来自定义过滤规则：

class Logger : public nvinfer1::ILogger { public: void log(Severity severity, const char* msg) noexcept override { if (severity <= Severity::kWARNING) { std::cout << "[TRT] " << msg << std::endl; } } } gLogger;

注意：这个gLogger必须是全局变量或静态变量，否则可能在推理过程中被析构，导致段错误。

核心类定义：YoloTRT

class YoloTRT { public: YoloTRT(const std::string& engine_path); ~YoloTRT(); float infer(cv::Mat& image, std::vector<Object>& result); private: void preprocess(const cv::Mat& input, cv::Mat& output); float* blobFromImage(cv::Mat& img); nvinfer1::ICudaEngine* engine; nvinfer1::IExecutionContext* context; cudaStream_t stream; void* buffers[5]; // 输入输出缓冲区 int inH, inW; size_t input_size, num_size, box_size, score_size, class_size; };

其中几个关键点值得强调：
-buffers数组保存GPU上各binding的指针地址，顺序需严格对应ONNX导出时的输出节点；
- 使用cudaStream_t实现异步传输与计算重叠；
-blobFromImage负责图像归一化（/255.0）、BGR→RGB转换以及NHWC→NCHW布局调整。

加载引擎：反序列化.trt文件

构造函数的核心任务是从磁盘加载已优化好的TensorRT引擎文件（.trt格式），并初始化运行时资源：

YoloTRT::YoloTRT(const std::string& engine_path) { std::ifstream file(engine_path, std::ios::binary | std::ios::ate); if (!file.is_open()) { std::cerr << "Cannot open engine file: " << engine_path << std::endl; abort(); } std::streamsize size = file.tellg(); std::vector<char> buffer(size); file.seekg(0, std::ios::beg); file.read(buffer.data(), size); file.close(); nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(gLogger); initLibNvInferPlugins(&gLogger, ""); // 注册自定义插件（如 EfficientNMS_TRT） engine = runtime->deserializeCudaEngine(buffer.data(), size); context = engine->createExecutionContext(); auto input_dims = engine->getBindingDimensions(0); inH = input_dims.d[2]; inW = input_dims.d[3]; num_size = engine->getBindingDimensions(1).d[1]; box_size = engine->getBindingDimensions(2).d[1] * engine->getBindingDimensions(2).d[2]; score_size = engine->getBindingDimensions(3).d[1]; class_size = engine->getBindingDimensions(4).d[1]; cudaMalloc(&buffers[0], input_size * sizeof(float)); cudaMalloc(&buffers[1], num_size * sizeof(int)); cudaMalloc(&buffers[2], box_size * sizeof(float)); cudaMalloc(&buffers[3], score_size * sizeof(float)); cudaMalloc(&buffers[4], class_size * sizeof(int)); cudaStreamCreate(&stream); runtime->destroy(); // runtime仅用于反序列化，之后可安全释放 }

这里有三点极易踩坑：
1.必须调用initLibNvInferPlugins：许多YOLO模型（尤其是YOLOX）导出时包含了自定义NMS插件，不注册会导致找不到kernel；
2.binding索引要对齐：建议用Netron打开ONNX模型确认输出节点顺序；
3.不要忘记销毁临时runtime对象：否则会浪费显存。

图像预处理：LetterBox缩放策略

为了保持原始宽高比、避免物体形变，采用letterbox方式进行resize。短边按比例缩放，长边两侧填充灰值（通常为114）：

void YoloTRT::preprocess(const cv::Mat& input, cv::Mat& output) { float r = std::min((float)inW / input.cols, (float)inH / input.rows); int unpad_w = r * input.cols; int unpad_h = r * input.rows; cv::Mat resized; cv::resize(input, resized, cv::Size(unpad_w, unpad_h), 0, 0, cv::INTER_LINEAR); int dw = inW - unpad_w; int dh = inH - unpad_h; dw /= 2; dh /= 2; cv::copyMakeBorder(resized, output, dh, dh, dw, dw, cv::BORDER_CONSTANT, cv::Scalar(114, 114, 114)); }

返回的比例因子r可用于后续将检测框坐标映射回原图空间。

主推理函数：异步执行与流水线优化

float YoloTRT::infer(cv::Mat& image, std::vector<Object>& result) { cv::Mat pr_img; preprocess(image, pr_img); float* host_input = blobFromImage(pr_img); cudaMemcpyAsync(buffers[0], host_input, input_size * sizeof(float), cudaMemcpyHostToDevice, stream); auto start = std::chrono::high_resolution_clock::now(); context->enqueueV2(buffers, stream, nullptr); int num_det; float* det_boxes = new float[box_size]; float* det_scores = new float[score_size]; int* det_classes = new int[class_size]; cudaMemcpyAsync(&num_det, buffers[1], sizeof(int), cudaMemcpyDeviceToHost, stream); cudaMemcpyAsync(det_boxes, buffers[2], box_size*sizeof(float), cudaMemcpyDeviceToHost, stream); cudaMemcpyAsync(det_scores, buffers[3], score_size*sizeof(float),cudaMemcpyDeviceToHost, stream); cudaMemcpyAsync(det_classes, buffers[4], class_size*sizeof(int), cudaMemcpyDeviceToHost, stream); cudaStreamSynchronize(stream); auto end = std::chrono::high_resolution_clock::now(); float scale = std::min((float)image.cols/inW, (float)image.rows/inH); int x_pad = (inW - image.cols/scale)/2; int y_pad = (inH - image.rows/scale)/2; for (int i = 0; i < num_det; ++i) { float score = det_scores[i]; if (score < 0.25) continue; float x0 = (det_boxes[i*4+0] - x_pad) * scale; float y0 = (det_boxes[i*4+1] - y_pad) * scale; float x1 = (det_boxes[i*4+2] - x_pad) * scale; float y1 = (det_boxes[i*4+3] - y_pad) * scale; Object obj; obj.rect = cv::Rect_<float>(x0, y0, x1-x0, y1-y0); obj.classId = det_classes[i]; obj.prob = score; result.push_back(obj); } delete[] host_input; delete[] det_boxes; delete[] det_scores; delete[] det_classes; return std::chrono::duration<float, std::milli>(end - start).count(); }

这里有几个性能优化技巧：
- 使用cudaMemcpyAsync而非同步拷贝，允许DMA传输与GPU计算并行；
- 所有输出拷贝操作都提交到同一个stream中，最后统一synchronize，减少同步开销；
- 检测阈值（如0.25）可根据实际场景动态调整，平衡速度与召回率。

资源清理：别让内存泄漏毁掉高性能

析构函数必须严格按照逆序销毁资源，尤其要注意CUDA流和显存的释放顺序：

YoloTRT::~YoloTRT() { cudaStreamSynchronize(stream); for (int i = 0; i < 5; ++i) cudaFree(buffers[i]); cudaStreamDestroy(stream); context->destroy(); engine->destroy(); }

任何一步遗漏都可能导致程序退出时报错或显存无法释放。

完整调用示例

struct Object { cv::Rect_<float> rect; int classId; float prob; }; int main(int argc, char** argv) { if (argc != 3) { std::cerr << "Usage: " << argv[0] << " <engine_file> <image_file>" << std::endl; return -1; } YoloTRT detector(argv[1]); cv::Mat img = cv::imread(argv[2]); if (img.empty()) { std::cerr << "Load image failed!" << std::endl; return -1; } // 预热几次，排除首次加载延迟 for (int i = 0; i < 5; ++i) { std::vector<Object> temp; detector.infer(img, temp); } std::vector<Object> results; float time_ms = detector.infer(img, results); std::cout << "Inference time: " << time_ms << " ms" << std::endl; for (const auto& obj : results) { cv::rectangle(img, obj.rect, cv::Scalar(0,255,0), 2); cv::putText(img, std::to_string(obj.classId), cv::Point(obj.rect.x, obj.rect.y-5), cv::FONT_HERSHEY_SIMPLEX, 0.6, cv::Scalar(0,0,255), 2); } cv::imwrite("output.jpg", img); std::cout << "Saved result to output.jpg" << std::endl; return 0; }

预热环节不可省略——首次推理往往包含上下文初始化、内存分配等额外开销，不能代表真实性能。

编译配置：CMake管理依赖

编写简洁高效的CMakeLists.txt：

cmake_minimum_required(VERSION 3.18) project(yolo_trt LANGUAGES CXX CUDA) set(CMAKE_CXX_STANDARD 14) find_package(OpenCV REQUIRED) find_package(CUDA REQUIRED) include_directories(${OpenCV_INCLUDE_DIRS}) add_executable(infer main.cpp) target_link_libraries(infer ${OpenCV_LIBS} nvinfer) set_property(TARGET infer PROPERTY CUDA_SEPARABLE_COMPILATION ON)

编译步骤：

mkdir build && cd build cmake .. && make -j8

运行：

./infer ../models/yolov5s.trt ../images/demo.jpg

预期输出类似：

Inference time: 8.76 ms Saved result to output.jpg

工程进阶方向与常见问题

这套基础框架已在多个项目中稳定运行，但仍有不少优化空间：

• INT8量化：配合校准集生成低精度引擎，可在几乎无损精度的前提下提速30%-50%；
• 动态Batch支持：修改engine profile以适应可变batch size，提升服务器吞吐；
• 多实例并发：为每个线程创建独立的IExecutionContext，实现单引擎多请求并行；
• Zero-Copy输入：对接摄像头驱动的DMA buffer，避免CPU-GPU重复拷贝；
• 迁移到新API：逐步采用ITensor、volumes等现代接口替代旧版绑定方式。

常见问题排查清单

从PyTorch原型到TensorRT+C++部署，这条路径看似陡峭，却是算法工程师走向工程闭环的必经之路。掌握这套技能不仅能显著提升系统性能，更能加深对深度学习底层运行机制的理解。当你看到一个原本需要几十毫秒的模型在嵌入式设备上稳定跑进10ms以内时，那种掌控感是无可替代的。

未来我会继续分享关于动态shape支持、INT8量化实战、多模型串联Pipeline设计等进阶主题，欢迎关注交流。如有疑问或发现文中疏漏，也欢迎指出，共同进步。