基于TensorRT和CUDA核函数加速YOLOv8推理：从零实现高效前处理

1次阅读

没有评论

共计 2410 个字符，预计需要花费 7 分钟才能阅读完成。

在目标检测的实际部署中，YOLOv8 模型推理通常只占整个流程的 30%-50% 时间，而前处理（包括图像缩放、归一化、颜色空间转换等）往往会消耗大量 CPU 资源。传统基于 OpenCV 的预处理存在以下问题：

串行执行导致吞吐量受限
频繁的 CPU-GPU 数据传输增加延迟
多线程管理复杂且容易引入竞争条件

通过实测，在 1080P 图像输入下，仅 resize 操作就可能占用 15ms 以上的 CPU 时间，这对于实时检测系统是不可接受的。

对比常见的前处理方案：

纯 OpenCV 方案
优点：开发简单，API 丰富
缺点：无法利用 GPU 并行性
OpenCV + CUDA 模块
优点：部分操作已优化
缺点：灵活性差，仍需要 CPU 参与
自定义 CUDA 核函数
优点：完全并行化，零拷贝传输
缺点：开发复杂度高

对于固定预处理流程，自定义 CUDA 核函数能实现最佳性能。以 YOLOv8 为例，其标准预处理包含：

BGR → RGB 转换
图像缩放 + 填充（letterbox）
归一化（/255.0）
NHWC → NCHW 转换

这些操作都适合用 CUDA 并行实现。

构建包含预处理阶段的 TensorRT 引擎：

// 创建 builder 和 network
auto builder = SampleUniquePtr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(logger));
auto network = SampleUniquePtr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(0));

// 添加输入层（接收原始图像数据）auto input = network->addInput("input", nvinfer1::DataType::kUINT8, 
    nvinfer1::Dims4{batch_size, height, width, channels});

// 添加预处理插件层（通过 CUDA 核函数实现）auto preprocess_layer = addPreprocessPlugin(*network, *input);

// 构建引擎
builder->setMaxBatchSize(batch_size);
auto config = builder->createBuilderConfig();
config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1 << 30);
auto engine = builder->buildEngineWithConfig(*network, *config);

实现高效的预处理核函数需要考虑：

内存访问模式
使用 coalesced memory access 确保全局内存合并访问
通过 shared memory 减少重复读取
线程组织
每个线程处理一个像素点（或一组像素）
blockDim.x 设置为 32 的倍数（一个 warp 的线程数）
数据搬运
使用 pinned memory 加速主机 - 设备传输
实现 zero-copy（如 CUDA Unified Memory）

示例核函数片段：

__global__ void preprocess_kernel(
    uint8_t* src, float* dst, 
    int in_w, int in_h, int out_size) {
    // 计算当前线程处理的像素位置
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    if (x >= in_w || y >= in_h) return;

    // 计算内存偏移（考虑 NHWC 布局）int src_offset = y * in_w * 3 + x * 3;
    int dst_offset = ...; // NCHW 布局计算

    // 执行实际预处理
    dst[dst_offset + 0] = src[src_offset + 2] / 255.0f; // R 通道
    dst[dst_offset + 1] = src[src_offset + 1] / 255.0f; // G 通道
    dst[dst_offset + 2] = src[src_offset + 0] / 255.0f; // B 通道
}

测试环境：
– GPU: RTX 3090
– CPU: AMD Ryzen 9 5950X
– 输入分辨率: 1920×1080
– 输出分辨率: 640×640