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

1次阅读

没有评论

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

在目标检测模型的部署中，前处理阶段往往成为性能瓶颈。传统 CPU 预处理（如 OpenCV 的 resize 和 normalize 操作）存在以下问题：

数据在 CPU 和 GPU 之间频繁搬运，导致 PCIe 带宽成为瓶颈
单线程处理无法充分利用现代多核 CPU 的算力
与 TensorRT 推理引擎异步执行困难，增加端到端延迟

以 YOLOv8 为例，输入图像通常需要经过以下处理链：

从原始字节流解码（JPEG/PNG 等）
调整尺寸至 640×640（保持长宽比填充灰边）
像素值归一化到 0 - 1 范围
转换为 NCHW 格式张量

实测表明，在 4K 输入分辨率下，仅 resize 操作就可能消耗 15ms 以上，而 TensorRT 推理本身仅需 8ms（RTX 3090）。这种比例失衡在实时系统中尤为致命。

常见前处理加速方案对比：

方案	延迟(ms)	GPU 利用率	实现复杂度	适用场景
OpenCV(cpu)	15.2	0%	低	原型验证
OpenCV(cuda)	5.8	30%	中	快速迭代
CUDA 核函数	1.2	85%	高	生产环境
DALI(pipeline)	3.5	60%	中高	视频流处理

CUDA 核函数的优势在于：

完全避免 CPU-GPU 数据传输
可精细控制线程调度和内存访问模式
与 TensorRT 推理无缝衔接（共享 CUDA stream）

关键配置参数：

const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
auto network = builder->createNetworkV2(explicitBatch);

// 配置优化 profile
auto profile = builder->createOptimizationProfile();
profile->setDimensions("images", OptProfileSelector::kMIN, Dims4{1, 3, 640, 640});
profile->setDimensions("images", OptProfileSelector::kOPT, Dims4{8, 3, 640, 640});
profile->setDimensions("images", OptProfileSelector::kMAX, Dims4{32, 3, 640, 640});

// 启用 FP16 和 TF32 加速
config->setFlag(BuilderFlag::kFP16);
config->setFlag(BuilderFlag::kTF32);

实现归一化和 resize 的融合操作：

__global__ void preprocess_kernel(
    uint8_t* src, float* dst, 
    int src_width, int src_height,
    int dst_width, int dst_height,
    float scale_x, float scale_y,
    float mean[3], float std[3]) {

    // 计算当前线程处理的输出坐标
    int dx = blockIdx.x * blockDim.x + threadIdx.x;
    int dy = blockIdx.y * blockDim.y + threadIdx.y;
    int dz = blockIdx.z;

    if (dx >= dst_width || dy >= dst_height || dz >= 3) return;

    // 计算对应输入坐标（双线性插值）float sx = (dx + 0.5f) * scale_x - 0.5f;
    float sy = (dy + 0.5f) * scale_y - 0.5f;

    int x0 = static_cast<int>(sx);
    int y0 = static_cast<int>(sy);
    int x1 = min(x0 + 1, src_width - 1);
    int y1 = min(y0 + 1, src_height - 1);

    // 计算插值权重
    float wx = sx - x0;
    float wy = sy - y0;

    // 读取四个相邻像素（coalesced 访问）uchar4 p00 = *reinterpret_cast<uchar4*>(src + (y0 * src_width + x0) * 3);
    uchar4 p01 = *reinterpret_cast<uchar4*>(src + (y0 * src_width + x1) * 3);
    uchar4 p10 = *reinterpret_cast<uchar4*>(src + (y1 * src_width + x0) * 3);
    uchar4 p11 = *reinterpret_cast<uchar4*>(src + (y1 * src_width + x1) * 3);

    // 通道分离和归一化
    float val = (1-wx)*(1-wy)*p00.x + wx*(1-wy)*p01.x + 
                (1-wx)*wy*p10.x + wx*wy*p11.x;
    val = (val / 255.0f - mean[dz]) / std[dz];

    // NCHW 布局写入
    int out_idx = dz * dst_width * dst_height + dy * dst_width + dx;
    dst[out_idx] = val;
}

void inference_pipeline(const std::vector<cv::Mat>& images) {
    // 1. 分配 pinned memory 输入缓冲区
    cudaMallocHost(&input_host, batch_size * 3 * 640 * 640 * sizeof(float));

    // 2. 启动 CUDA 核函数
    dim3 block(32, 32);
    dim3 grid((640 + block.x - 1) / block.x,
        (640 + block.y - 1) / block.y,
        3);

    preprocess_kernel<<<grid, block, 0, stream>>>(...);

    // 3. TensorRT 异步推理
    void* bindings[] = {input_device, output_device};
    context->enqueueV2(bindings, stream, nullptr);

    // 4. 后处理（略）}

测试环境：RTX 3090, CUDA 11.7, TensorRT 8.4