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

在边缘计算和实时视觉系统日益普及的今天，如何让YOLO这类高效目标检测模型真正“跑得快、稳得住”，成了从算法到落地的关键一环。项目中我们常遇到这样的场景：PyTorch训练出的模型精度不错，转ONNX后用Python推理也能出图，但一旦进入生产环境——延迟高、吞吐低、GPU利用率上不去，系统瓶颈立刻显现。

于是，转向 TensorRT + C++ 成为必然选择。这不是为了炫技，而是贴近工业级部署的真实需求：自动驾驶模块要毫秒级响应，智能摄像头需7×24小时运行，视频分析平台得处理上百路流。这些场景下，C++带来的不仅是性能提升，更是一种工程思维的转变——对内存、异步、资源释放的精确控制。

本文记录我在Linux服务器上基于NVIDIA官方镜像完成YOLOv5、YOLOX等系列模型C++推理部署的全过程。不讲理论堆砌，只聚焦实战细节：环境怎么搭、代码怎么写、常见坑怎么避。希望正在从Python转型工程落地的朋友能少踩几个坑。

---

环境准备：用Docker构建可复现的推理容器

实验运行于一台双NVIDIA 3090显卡的Ubuntu 20.04服务器，采用Docker容器化方式隔离依赖，确保环境纯净且易于迁移。

使用NVIDIA官方TensorRT镜像

NVIDIA提供了预装CUDA、cuDNN、TensorRT及相关工具链的官方镜像，极大简化了配置复杂度。这次选用的是 nvcr.io/nvidia/tensorrt:23.10-py3，其内置：

• CUDA 12.2
• TensorRT 8.6.1
• cuDNN 8.9
• 支持FP16/INT8量化
• 已编译好的库文件与头文件

启动命令如下：

sudo docker run -it \
  -v /home/user/project:/workspace \
  --workdir=/workspace \
  --network=host \
  --runtime=nvidia \
  --gpus all \
  --shm-size=16g \
  --name trt_yolo \
  nvcr.io/nvidia/tensorrt:23.10-py3 /bin/bash

⚠️ 注意事项：
- 宿主机需安装对应版本驱动（建议 >= 535）
- 若驱动较低，请选兼容旧版镜像（如22.xx）

安装OpenCV及其他依赖

虽然TensorRT已就位，但图像处理仍需OpenCV。推荐使用APT快速安装，避免源码编译带来的路径混乱问题：

apt update && apt upgrade -y

# 编译工具链
apt install -y build-essential g++ cmake make

# OpenCV开发包及多媒体支持
apt install -y libopencv-dev python3-dev python3-numpy \
               libavcodec-dev libavformat-dev libswscale-dev \
               libgtk2.0-dev libgtk-3-dev libpng-dev libtiff-dev \
               libwebp-dev libopenexr-dev libjpeg-dev

# 验证安装
pkg-config --modversion opencv4  # 应输出4.x版本号

至此，推理所需的基础环境已齐备。接下来就可以专注写核心逻辑了。

---

核心设计：两个类撑起整个推理流程

整个C++推理架构围绕两个关键类展开：Logger 和 Yolo。结构简洁，却覆盖了从日志管理到端到端推理的完整生命周期。

Logger类：接管TensorRT的日志输出

所有TensorRT API调用都会通过 nvinfer1::ILogger 接口打印日志。默认行为是全量输出，信息过于冗长。我们需要自定义一个子类来过滤掉无用信息。

class Logger : public nvinfer1::ILogger {
public:
    void log(Severity severity, const char* msg) noexcept override {
        if (severity != Severity::kINFO) {  // 只保留warn/error级别
            std::cout << "[TRT] " << msg << std::endl;
        }
    }
};

这个轻量实现可以显著减少控制台干扰。实际调试时，你可能会临时放开 kINFO 查看层融合详情；上线后则建议只留 kERROR 或 kINTERNAL_ERROR，保持日志干净。

该 logger 实例将在后续创建 IRuntime 和反序列化引擎时传入，是几乎所有操作的前提。

---

Yolo类：封装模型加载与推理主干

Yolo 类承担了从 .engine 文件加载、GPU内存分配、预处理到后处理的全流程任务。

类声明概览

class Yolo {
public:
    Yolo(const char* engine_file_path);
    ~Yolo();

    float letterbox(const cv::Mat& image, cv::Mat& resized, cv::Size new_shape = {640, 640});
    float* blobFromImage(cv::Mat& img);
    void draw_objects(cv::Mat& img, float* boxes, int* classes, int num_boxes);

    void Infer(
        int width, int height, int channel,
        unsigned char* input_data,
        float* output_boxes,
        int* output_classes,
        int* num_boxes
    );

private:
    nvinfer1::IRuntime* runtime;
    nvinfer1::ICudaEngine* engine;
    nvinfer1::IExecutionContext* context;
    cudaStream_t stream;

    void* buffers[5];  // 输入输出缓冲区指针
    int iH, iW;        // 模型输入高度和宽度
    size_t in_size, out_size1, out_size2, out_size3, out_size4;

    Logger logger;
};

其中 buffers[5] 是五个绑定张量的统一指针数组，顺序需与 .engine 中binding索引一致（通常为：images → num_dets → det_boxes → det_scores → det_classes）。

---

构造函数：重建推理上下文

构造函数负责读取 .engine 文件并重建完整的推理环境。

Yolo::Yolo(const char* engine_file_path) : runtime(nullptr), engine(nullptr), context(nullptr) {
    std::ifstream file(engine_file_path, std::ios::binary | std::ios::ate);
    if (!file.is_open()) {
        std::cerr << "Cannot open engine file: " << engine_file_path << std::endl;
        abort();
    }

    size_t fileSize = file.tellg();
    file.seekg(0, std::ios::beg);
    std::vector<char> engineData(fileSize);
    file.read(engineData.data(), fileSize);
    file.close();

    runtime = nvinfer1::createInferRuntime(logger);
    if (!runtime) {
        std::cerr << "Failed to create TensorRT runtime." << std::endl;
        abort();
    }

    initLibNvInferPlugins(&logger, "");  // 注册插件（如NMS）

    engine = runtime->deserializeCudaEngine(engineData.data(), fileSize);
    if (!engine) {
        std::cerr << "Failed to deserialize engine from buffer." << std::endl;
        abort();
    }

    context = engine->createExecutionContext();
    if (!context) {
        std::cerr << "Failed to create execution context." << std::endl;
        abort();
    }

    auto input_dims = engine->getBindingDimensions(engine->getBindingIndex("images"));
    iH = input_dims.d[2];
    iW = input_dims.d[3];

    in_size = 1;
    for (int i = 0; i < input_dims.nbDims; ++i)
        in_size *= input_dims.d[i];

    out_size1 = getSizeByDim(engine->getBindingDimensions(engine->getBindingIndex("num_dets")));
    out_size2 = getSizeByDim(engine->getBindingDimensions(engine->getBindingIndex("det_boxes")));
    out_size3 = getSizeByDim(engine->getBindingDimensions(engine->getBindingIndex("det_scores")));
    out_size4 = getSizeByDim(engine->getBindingDimensions(engine->getBindingIndex("det_classes")));

    cudaMalloc(&buffers[0], in_size  * sizeof(float));
    cudaMalloc(&buffers[1], out_size1 * sizeof(int));
    cudaMalloc(&buffers[2], out_size2 * sizeof(float));
    cudaMalloc(&buffers[3], out_size3 * sizeof(float));
    cudaMalloc(&buffers[4], out_size4 * sizeof(int));

    cudaStreamCreate(&stream);
}

辅助函数 getSizeByDim() 计算维度乘积：

size_t getSizeByDim(const nvinfer1::Dims& dims) {
    size_t size = 1;
    for (int i = 0; i < dims.nbDims; ++i)
        size *= dims.d[i];
    return size;
}

💡 经验提示：若你的模型有多个输出分支或动态shape，建议将 out_sizeX 存入数组而非硬编码变量，便于扩展。

---

letterbox：保持宽高比的缩放策略

直接拉伸图像会导致形变，影响小物体检测效果。因此采用letterbox方式，在缩放后填充灰边（值为114，符合YOLO训练时的数据增强设定）。

float Yolo::letterbox(
    const cv::Mat& image,
    cv::Mat& resized,
    cv::Size new_shape
) {
    float r = std::min((float)new_shape.height / image.rows, (float)new_shape.width / image.cols);
    int unpad_w = (int)(image.cols * r);
    int unpad_h = (int)(image.rows * r);

    cv::Mat tmp;
    cv::resize(image, tmp, cv::Size(unpad_w, unpad_h));

    int dw = new_shape.width - unpad_w;
    int dh = new_shape.height - unpad_h;
    dw /= 2;
    dh /= 2;

    cv::copyMakeBorder(tmp, resized, 
                       (int)std::round(dh - 0.1f), (int)std::round(dh + 0.1f),
                       (int)std::round(dw - 0.1f), (int)std::round(dw + 0.1f),
                       cv::BORDER_CONSTANT, cv::Scalar(114, 114, 114));

    return 1.f / r;  // 返回逆缩放因子，用于坐标还原
}

返回 1/r 是为了后续将网络输出的bbox映射回原始图像空间，这点很容易被忽略但至关重要。

---

blobFromImage：图像归一化与格式转换

将 cv::Mat 转换为模型所需的浮点输入张量，并执行归一化（除以255）和BGR→RGB通道反转。

float* Yolo::blobFromImage(cv::Mat& img) {
    float* blob = new float[in_size];
    int channels = 3;
    int img_area = img.cols * img.rows;
    uchar* pimg = img.data;

    for (int c = 0; c < channels; ++c) {
        for (int i = 0; i < img_area; ++i) {
            blob[c * img_area + i] = (float)pimg[i * channels + (2 - c)] / 255.0f;
        }
    }
    return blob;
}

🔍 注意：(2 - c) 实现了从BGR到RGB的反转。如果你的输入已经是RGB格式（如某些相机SDK直接输出），这里应改为 (c)。

---

Infer函数：五步走完一次推理

这是最核心的部分，遵循标准的TensorRT推理流程：

1. 图像预处理（letterbox + RGB转换）
2. Host → Device 异步拷贝
3. 执行 enqueueV2 启动推理
4. Device → Host 异步回传结果
5. 坐标还原 + 边界裁剪

void Yolo::Infer(
    int width, int height, int channel,
    unsigned char* input_data,
    float* output_boxes,
    int* output_classes,
    int* num_boxes
) {
    cv::Mat src(height, width, CV_8UC3, input_data);
    cv::Mat pr_img;
    float scale = letterbox(src, pr_img, {iW, iH});

    cv::cvtColor(pr_img, pr_img, cv::COLOR_BGR2RGB);
    float* blob = blobFromImage(pr_img);

    static int   det_num;
    static float det_boxes[4000 * 4];
    static float det_scores[4000];
    static int   det_classes[4000];

    cudaMemcpyAsync(buffers[0], blob, in_size * sizeof(float), cudaMemcpyHostToDevice, stream);
    context->enqueueV2(buffers, stream, nullptr);
    cudaMemcpyAsync(&det_num,     buffers[1], out_size1 * sizeof(int),   cudaMemcpyDeviceToHost, stream);
    cudaMemcpyAsync(det_boxes,    buffers[2], out_size2 * sizeof(float), cudaMemcpyDeviceToHost, stream);
    cudaMemcpyAsync(det_scores,   buffers[3], out_size3 * sizeof(float), cudaMemcpyDeviceToHost, stream);
    cudaMemcpyAsync(det_classes, buffers[4], out_size4 * sizeof(int),   cudaMemcpyDeviceToHost, stream);

    cudaStreamSynchronize(stream);  // 等待所有操作完成

    int x_offset = (iW * scale - src.cols) / 2;
    int y_offset = (iH * scale - src.rows) / 2;

    num_boxes[0] = det_num;
    for (int i = 0; i < det_num; ++i) {
        float x0 = det_boxes[i * 4 + 0] * scale - x_offset;
        float y0 = det_boxes[i * 4 + 1] * scale - y_offset;
        float x1 = det_boxes[i * 4 + 2] * scale - x_offset;
        float y1 = det_boxes[i * 4 + 3] * scale - y_offset;

        x0 = std::max(std::min(x0, (float)(src.cols - 1)), 0.f);
        y0 = std::max(std::min(y0, (float)(src.rows - 1)), 0.f);
        x1 = std::max(std::min(x1, (float)(src.cols - 1)), 0.f);
        y1 = std::max(std::min(y1, (float)(src.rows - 1)), 0.f);

        output_boxes[i * 4 + 0] = x0;
        output_boxes[i * 4 + 1] = y0;
        output_boxes[i * 4 + 2] = x1 - x0;
        output_boxes[i * 4 + 3] = y1 - y0;
        output_classes[i] = det_classes[i];
    }

    delete[] blob;
}

🛠️ 工程建议：
- static 数组用于避免频繁new/delete，适合固定最大检测数场景
- 若需支持动态数量，可改用 std::vector<float> 并预留容量
- 坐标还原部分加入边界裁剪，防止越界访问导致崩溃

---

析构函数：别忘了释放GPU资源

C++没有GC，所有手动申请的资源都必须显式释放，否则极易造成内存泄漏甚至程序崩溃。

Yolo::~Yolo() {
    cudaStreamSynchronize(stream);

    cudaFree(buffers[0]);
    cudaFree(buffers[1]);
    cudaFree(buffers[2]);
    cudaFree(buffers[3]);
    cudaFree(buffers[4]);

    cudaStreamDestroy(stream);
    if (context) context->destroy();
    if (engine)  engine->destroy();
    if (runtime) runtime->destroy();
}

尤其注意：destroy() 是TensorRT对象特有的析构方法，不能用 delete 替代。

---

主函数示例：完整调用流程

int main(int argc, char** argv) {
    if (argc != 5 ||
        std::string(argv[1]) != "-model_path" ||
        std::string(argv[3]) != "-image_path") {
        std::cerr << "Usage: " << argv[0] 
                  << " -model_path <path_to_engine> -image_path <path_to_image>" 
                  << std::endl;
        return -1;
    }

    const char* model_path = argv[2];
    const char* image_path = argv[4];

    Yolo yolo(model_path);
    cv::Mat img = cv::imread(image_path);
    if (img.empty()) {
        std::cerr << "Failed to load image: " << image_path << std::endl;
        return -1;
    }

    float* boxes = new float[4000 * 4];
    int* classes = new int[4000];
    int* num_boxes = new int[1];

    // 预热几次，排除首次调度开销
    for (int i = 0; i < 10; ++i) {
        yolo.Infer(img.cols, img.rows, img.channels(), img.data, boxes, classes, num_boxes);
    }

    auto start = std::chrono::steady_clock::now();
    yolo.Infer(img.cols, img.rows, img.channels(), img.data, boxes, classes, num_boxes);
    auto end = std::chrono::steady_clock::now();

    auto ms = std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.0;
    std::cout << "Inference time: " << ms << " ms" << std::endl;

    yolo.draw_objects(img, boxes, classes, num_boxes);
    cv::imwrite("result.jpg", img);

    delete[] boxes;
    delete[] classes;
    delete[] num_boxes;

    return 0;
}

预热（warm-up）非常重要！首次推理往往包含上下文初始化、内核加载等额外开销，实测可能高出正常值30%以上。

---

性能表现与进阶优化方向

在我的3090设备上测试YOLOv5s-engine（FP16精度），单图推理时间稳定在 ~4.2ms，即约 238 FPS，相较PyTorch原生推理提速近5倍。

这还只是起点。以下几个方向可进一步压榨性能：

优化手段	效果
FP16 精度推理	提速30%-50%，显存减半
INT8 量化校准	再提速20%-40%，几乎无精度损失
多batch推理	提高GPU利用率，适合视频流批处理
动态profile适配不同输入尺寸	更灵活应对多分辨率输入

特别是INT8量化，在精度损失小于1%的情况下，常常能再带来30%以上的速度提升。配合校准集生成scale文件，是工业部署中的标配操作。

此外，结合 DeepStream SDK 可轻松实现多路视频流并行处理，适用于安防监控、交通卡口等大规模场景。

---

这种高度集成的C++推理方案，正成为智能视觉系统的底层支柱。它不像Jupyter Notebook那样直观易调，但一旦掌握其模式——内存管理、异步流同步、RAII资源守恒——就能真正释放GPU的极限性能。未来我还会深入探索动态shape支持、自定义Plugin开发、ONNX Parser自动化建图等方向，持续打磨这套“看得清、跑得快”的视觉引擎。