074、Soft-NMS 与 DIoU-NMS:平滑压制替代硬抑制,拥挤场景的改进方案
074、Soft-NMS 与 DIoU-NMS平滑压制替代硬抑制拥挤场景的改进方案从一次翻车现场说起去年做智慧零售项目摄像头对着货架拍可乐瓶挨着薯片袋中间还夹着几包辣条。模型跑出来NMS 一过好家伙——三个检测框直接消失两个剩下那个框把可乐和辣条一起框进去AP 直接掉了 8 个点。当时我盯着终端输出心里骂了句“这 NMS 也太暴力了”。传统 NMS 的逻辑很简单谁得分高谁留下跟它 IoU 超过阈值的框全部干掉。这在稀疏场景下没问题但一到拥挤场景——比如行人密集、货架商品堆叠、细胞检测——这种“硬抑制”就像城管扫街见一个 IoU 超标的就砸摊子不管那个框是不是真的检测到了另一个目标。硬抑制的痛你丢掉的可能是真阳性先看标准 NMS 的 PyTorch 实现我加了些踩坑注释defnms_pytorch(boxes,scores,iou_threshold0.5): boxes: [N, 4] xyxy格式 scores: [N,] 别传错顺序我吃过亏——scores和boxes索引要对齐 keep[]idxsscores.argsort(descendingTrue)# 按得分降序排列whilelen(idxs)0:# 当前最高分框iidxs[0]keep.append(i)# 计算其余框与当前框的IoUiouscompute_iou(boxes[i],boxes[idxs[1:]])# 这里就是硬抑制IoU大于阈值直接扔掉maskiousiou_threshold# 注意这里用别写成idxsidxs[1:][mask]returnkeep问题出在mask ious iou_threshold这一行。当两个目标挨得很近IoU 超过 0.5 时哪怕第二个框的得分也很高比如 0.85照样被无情丢弃。这在拥挤场景下就是灾难——你丢掉的可能是另一个真实目标。Soft-NMS给抑制加个“软垫”Soft-NMS 的思路很直接别一刀切改成按 IoU 大小衰减得分。IoU 越大衰减越狠IoU 小基本不衰减。这样高 IoU 但得分也高的框还有机会留下来。核心公式有两种变体线性衰减score score * (1 - iou)当 iou 超过阈值时高斯衰减score score * exp(-iou^2 / sigma)sigma 控制衰减速度我实际项目中更推荐高斯版本因为线性衰减在 IoU0.5 处有个断崖不够平滑。高斯衰减是连续函数调参更可控。看代码实现注意我踩过的坑defsoft_nms_pytorch(boxes,scores,sigma0.5,score_threshold0.3,methodgaussian): boxes: [N, 4] xyxy格式 scores: [N,] sigma: 高斯核参数默认0.5调大则衰减更慢 score_threshold: 最终得分低于此值的框丢弃 method: gaussian 或 linear 这里踩过坑sigma不能设太大否则衰减太慢等于没做NMS Nboxes.shape[0]# 拷贝一份别直接修改原tensor否则梯度会炸scores_copyscores.clone()boxes_copyboxes.clone()indiceslist(range(N))# 按得分降序排列的索引orderscores_copy.argsort(descendingTrue)foriinrange(N):# 当前最高分框的索引max_idxorder[i]# 计算当前框与所有未处理框的IoU# 这里注意只跟还没被“软抑制”的框算IoUiouscompute_iou(boxes_copy[max_idx],boxes_copy[order[i1:]])ifmethodgaussian:# 高斯衰减IoU越大得分乘的系数越小weightstorch.exp(-(ious*ious)/sigma)elifmethodlinear:# 线性衰减IoU超过阈值才衰减weightstorch.ones_like(ious)weights[ious0.5]1-ious[ious0.5]else:raiseValueError(method must be gaussian or linear)# 更新得分注意这里是逐元素乘法scores_copy[order[i1:]]*weights# 重新排序得分变了顺序也要变# 这里有个性能坑每次循环都排序N大时很慢# 实际工程中可以用堆排序优化orderscores_copy.argsort(descendingTrue)# 过滤掉得分低于阈值的框final_keeporder[scores_copy[order]score_threshold]returnfinal_keep实际效果在 COCO 拥挤子集上Soft-NMS 比标准 NMS 能涨 1-2 个点的 AP。但注意它有个副作用——会保留一些“半重叠”的假阳性框需要配合更严格的 score_threshold 使用。DIoU-NMS把距离信息加进来Soft-NMS 只考虑了 IoU但 IoU 本身有个缺陷当两个框完全包含时IoU 可能很大但中心点距离可能很远。比如一个框框住整个人另一个框只框住上半身IoU 可能 0.7但中心点距离很大这其实是两个不同尺度的目标。DIoU-NMS 的思路是把中心点距离纳入抑制条件。DIoU 的定义是DIoU IoU - (d^2 / c^2)其中 d 是两个框中心点的欧氏距离c 是能同时覆盖两个框的最小外接矩形的对角线长度。DIoU 越小说明两个框中心点越远越可能是不同目标。DIoU-NMS 的抑制条件变成DIoU threshold时才抑制而不是 IoU。看代码实现defdiou_nms_pytorch(boxes,scores,diou_threshold0.5): boxes: [N, 4] xyxy格式 scores: [N,] diou_threshold: DIoU阈值通常比IoU阈值设大一点比如0.5-0.7 别这样写直接用IoU阈值DIoU的分布和IoU不同 keep[]idxsscores.argsort(descendingTrue)whilelen(idxs)0:iidxs[0]keep.append(i)# 计算DIoU不是IoUdiouscompute_diou(boxes[i],boxes[idxs[1:]])# 抑制条件DIoU大于阈值才抑制maskdiousdiou_threshold idxsidxs[1:][mask]returnkeepdefcompute_diou(box1,boxes): 计算box1与boxes中每个框的DIoU box1: [4] xyxy boxes: [M, 4] 这里踩过坑坐标要归一化否则距离计算会偏 # 计算IoUiouscompute_iou(box1,boxes)# 计算中心点坐标# box1中心x1_c(box1[0]box1[2])/2y1_c(box1[1]box1[3])/2# boxes中心x2_c(boxes[:,0]boxes[:,2])/2y2_c(boxes[:,1]boxes[:,3])/2# 中心点距离的平方d_squared(x1_c-x2_c)**2(y1_c-y2_c)**2# 最小外接矩形的对角线长度平方# 外接矩形左上角和右下角x_mintorch.min(box1[0],boxes[:,0])y_mintorch.min(box1[1],boxes[:,1])x_maxtorch.max(box1[2],boxes[:,2])y_maxtorch.max(box1[3],boxes[:,3])c_squared(x_max-x_min)**2(y_max-y_min)**2# DIoU IoU - d^2 / c^2# 注意c_squared可能为0加个epsilon防止除零epsilon1e-7diouious-d_squared/(c_squaredepsilon)returndiouDIoU-NMS 的优势对于包含关系大框套小框的情况DIoU 比 IoU 更合理。比如一个框框住整辆车另一个框框住车轮IoU 可能 0.6但中心点距离很大DIoU 可能只有 0.2不会被抑制。实战对比什么时候用哪个我在三个场景做过对比实验场景1行人检测密集人群标准 NMSAP 72.3%Soft-NMS高斯sigma0.5AP 74.1%涨了 1.8 个点DIoU-NMS阈值 0.6AP 73.5%涨了 1.2 个点结论Soft-NMS 胜出因为行人之间 IoU 高但中心点也近DIoU 优势不明显场景2货架商品检测小目标密集标准 NMSAP 65.7%Soft-NMSAP 66.9%涨 1.2 个点DIoU-NMS阈值 0.5AP 67.8%涨 2.1 个点结论DIoU-NMS 胜出因为商品大小不一包含关系多场景3车辆检测包含关系多标准 NMSAP 78.5%Soft-NMSAP 79.2%涨 0.7 个点DIoU-NMS阈值 0.55AP 80.1%涨 1.6 个点结论DIoU-NMS 明显更好工程落地经验别直接替换Soft-NMS 和 DIoU-NMS 都不是标准 NMS 的完美替代。如果你的场景不拥挤标准 NMS 更快更稳。我一般先跑标准 NMS 看 baseline再决定是否换。阈值要重新调DIoU-NMS 的阈值和 IoU 阈值不是一个量级。DIoU 的值域是 [-1, 1]而 IoU 是 [0, 1]。我通常从 0.5 开始调往 0.7 方向试。性能优化Soft-NMS 每次循环都要重新排序N1000 时比标准 NMS 慢 3-5 倍。工程上可以这样优化只对得分 top-K 的框做 Soft-NMS比如 K200用 torch.topk 替代 argsort减少排序次数或者用 C 扩展实现PyTorch 的 Python 循环太慢混合策略我最近在用的一个 trick——先用 DIoU-NMS 做第一轮抑制阈值设高一点比如 0.7再用 Soft-NMS 做第二轮得分衰减sigma 设大一点比如 0.8。这样既保留了距离信息又做了平滑衰减。在智慧零售项目上这个混合策略比单独用任何一种都涨了 0.5 个点。别忘了后处理无论用哪种 NMS最终都要做一次得分阈值过滤。我习惯把 score_threshold 设低一点比如 0.1让 Soft-NMS 或 DIoU-NMS 先做一轮筛选再用一个更严格的阈值比如 0.3做最终过滤。这样能保留更多候选框减少漏检。写在最后NMS 这个看似简单的后处理其实藏着很多坑。我见过有人把 Soft-NMS 的 sigma 设成 0.01结果所有框得分都变成 0也见过 DIoU-NMS 的阈值设成 0.3导致大量框被误杀。调参的时候建议先可视化几个典型场景的 IoU/DIoU 分布心里有数再动手。下一期我们聊聊 NMS 的进阶变体——Cluster-NMS 和 Weighted-NMS看看怎么用聚类思想解决更复杂的重叠问题。