YOLO姿态估计批量推理中首次NMS迭代耗时异常排查

问题原因分析 #

1. PyTorch首次初始化开销 #

首次执行CUDA操作时，PyTorch会完成CUDA上下文初始化、CUDA kernel JIT编译，同时如果有CPU张量首次迁移到GPU，还会产生数据拷贝开销。这些操作仅在第一次执行时触发，直接导致首帧NMS耗时显著高于后续帧。

2. 逐帧NMS的低效性 #

如果当前non_max_suppression_kpt是逐batch元素循环处理，首帧循环内的张量索引、IoU计算等操作会触发kernel编译，后续帧复用已编译的kernel，因此速度更快。这种逐帧处理的方式本质上不是真正的批量NMS，无法充分利用GPU的并行计算能力。

排查方法 #

验证初始化开销 #

在正式推理前添加预热步骤：

# 构造与实际输入同形状的dummy数据
dummy_pred = torch.randn((5, 1000, 5 + 17*3), device='cuda')
# 跑一遍NMS预热
non_max_suppression_kpt(dummy_pred)
# 再执行正式推理
results = non_max_suppression_kpt(actual_pred)

如果预热后首帧耗时与后续帧接近，说明问题源于首次初始化。

验证逐帧NMS的瓶颈 #

对比逐帧处理和批量处理的耗时：如果批量处理后所有帧耗时一致，说明当前代码不支持批量NMS是核心问题。

批量NMS改造方案 #

针对YOLO姿态估计的non_max_suppression_kpt，改造思路是保留batch维度，统一处理所有batch的框，通过坐标偏移避免跨batch抑制，最终按batch分组输出。以下是改造后的代码：

import torch
import torchvision

def xywh2xyxy(x):
    # 转换框格式：xywh -> xyxy
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[..., 0] = x[..., 0] - x[..., 2] / 2  # top left x
    y[..., 1] = x[..., 1] - x[..., 3] / 2  # top left y
    y[..., 2] = x[..., 0] + x[..., 2] / 2  # bottom right x
    y[..., 3] = x[..., 1] + x[..., 3] / 2  # bottom right y
    return y

def non_max_suppression_kpt(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, max_det=300):
    batch_size = prediction.shape[0]
    # 计算类别数（假设每个关键点包含x/y/conf三个值，这里以17个关键点为例）
    num_kpts = 17
    nc = prediction.shape[2] - 5 - num_kpts * 3
    xc = prediction[..., 4] > conf_thres  # 置信度过滤

    # 收集所有batch的有效框、分数、关键点及batch索引
    all_boxes = []
    all_scores = []
    all_kpts = []
    batch_indices = []

    for xi in range(batch_size):
        x = prediction[xi][xc[xi]]
        if not x.shape[0]:
            continue
        # 计算最终置信度：目标置信度 × 类别置信度
        x[:, 5:5+nc] *= x[:, 4:5]
        # 转换框格式
        box = xywh2xyxy(x[:, :4])
        # 提取关键点数据
        kpts = x[:, 5+nc:]
        # 取最大类别的置信度作为框的分数
        scores = x[:, 5:5+nc].max(dim=1)[0]
        
        all_boxes.append(box)
        all_scores.append(scores)
        all_kpts.append(kpts)
        batch_indices.append(torch.full((x.shape[0],), xi, device=x.device))

    if not all_boxes:
        # 无有效框时返回空结果
        return [torch.zeros((0, 6 + num_kpts*3), device=prediction.device) for _ in range(batch_size)]

    # 合并所有batch的数据
    all_boxes = torch.cat(all_boxes, dim=0)
    all_scores = torch.cat(all_scores, dim=0)
    all_kpts = torch.cat(all_kpts, dim=0)
    batch_indices = torch.cat(batch_indices, dim=0)

    # 添加batch坐标偏移，避免跨batch的框被NMS抑制
    if not agnostic:
        # 计算每个batch的偏移量，确保跨batch框无重叠
        x_offset = batch_indices * (all_boxes[:, 2].max() + 10)
        y_offset = batch_indices * (all_boxes[:, 3].max() + 10)
        all_boxes[:, [0, 2]] += x_offset
        all_boxes[:, [1, 3]] += y_offset

    # 执行批量NMS
    idx = torchvision.ops.nms(all_boxes, all_scores, iou_thres)
    if idx.shape[0] > max_det:
        idx = idx[:max_det]

    # 恢复原始坐标
    if not agnostic:
        x_offset = batch_indices[idx] * (all_boxes[:, 2].max() + 10)
        y_offset = batch_indices[idx] * (all_boxes[:, 3].max() + 10)
        all_boxes[idx, [0, 2]] -= x_offset
        all_boxes[idx, [1, 3]] -= y_offset

    # 按batch索引分组结果
    output = [torch.zeros((0, 6 + num_kpts*3), device=prediction.device) for _ in range(batch_size)]
    for i in idx:
        b_idx = batch_indices[i].item()
        # 拼接框、分数、关键点
        res_item = torch.cat([all_boxes[i], all_scores[i:i+1], all_kpts[i]], dim=0)
        # 追加到对应batch的结果中
        if output[b_idx].shape[0] == 0:
            output[b_idx] = res_item.unsqueeze(0)
        else:
            output[b_idx] = torch.cat([output[b_idx], res_item.unsqueeze(0)], dim=0)

    return output

改造核心点 #

1. 批量收集有效框：不再逐帧处理，一次性收集所有batch的有效检测框、分数和关键点。
2. 跨batch隔离：通过给不同batch的框添加坐标偏移，确保NMS只在同一batch内的框之间进行。
3. 官方NMS调用：使用torchvision.ops.nms实现批量NMS，利用GPU并行计算，避免循环内的重复kernel编译。
4. 结果分组：处理完成后按batch索引分组，保持与原函数一致的输出格式。

额外优化建议 #

• 移除频繁的torch.cuda.empty_cache()调用：该操作会触发显存清理，首次执行开销极大，PyTorch会自动管理显存，仅在显存不足时调用即可。
• 预热整个推理流程：不仅预热NMS，还可以预热模型前向传播，进一步消除首帧初始化开销。

内容的提问来源于stack exchange，提问作者gamerchief gaming