当前位置：首页 > news >正文

074、Soft-NMS 与 DIoU-NMS：平滑压制替代硬抑制，拥挤场景的改进方案

news 2026/6/11 5:22:17

074、Soft-NMS 与 DIoU-NMS：平滑压制替代硬抑制，拥挤场景的改进方案

从一次翻车现场说起

去年做智慧零售项目，摄像头对着货架拍，可乐瓶挨着薯片袋，中间还夹着几包辣条。模型跑出来，NMS 一过，好家伙——三个检测框直接消失两个，剩下那个框把可乐和辣条一起框进去，AP 直接掉了 8 个点。当时我盯着终端输出，心里骂了句“这 NMS 也太暴力了”。

传统 NMS 的逻辑很简单：谁得分高谁留下，跟它 IoU 超过阈值的框全部干掉。这在稀疏场景下没问题，但一到拥挤场景——比如行人密集、货架商品堆叠、细胞检测——这种“硬抑制”就像城管扫街，见一个 IoU 超标的就砸摊子，不管那个框是不是真的检测到了另一个目标。

硬抑制的痛：你丢掉的可能是真阳性

先看标准 NMS 的 PyTorch 实现，我加了些踩坑注释：

defnms_pytorch(boxes,scores,iou_threshold=0.5):""" boxes: [N, 4] xyxy格式 scores: [N,] 别传错顺序，我吃过亏——scores和boxes索引要对齐 """keep=[]idxs=scores.argsort(descending=True)# 按得分降序排列whilelen(idxs)>0:# 当前最高分框i=idxs[0]keep.append(i)# 计算其余框与当前框的IoUious=compute_iou(boxes[i],boxes[idxs[1:]])# 这里就是硬抑制：IoU大于阈值直接扔掉mask=ious<=iou_threshold# 注意：这里用<=，别写成<idxs=idxs[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 控制衰减速度

我实际项目中更推荐高斯版本，因为线性衰减在 IoU=0.5 处有个断崖，不够平滑。高斯衰减是连续函数，调参更可控。

看代码实现，注意我踩过的坑：

defsoft_nms_pytorch(boxes,scores,sigma=0.5,score_threshold=0.3,method='gaussian'):""" boxes: [N, 4] xyxy格式 scores: [N,] sigma: 高斯核参数，默认0.5，调大则衰减更慢 score_threshold: 最终得分低于此值的框丢弃 method: 'gaussian' 或 'linear' 这里踩过坑：sigma不能设太大，否则衰减太慢等于没做NMS """N=boxes.shape[0]# 拷贝一份，别直接修改原tensor，否则梯度会炸scores_copy=scores.clone()boxes_copy=boxes.clone()indices=list(range(N))# 按得分降序排列的索引order=scores_copy.argsort(descending=True)foriinrange(N):# 当前最高分框的索引max_idx=order[i]# 计算当前框与所有未处理框的IoU# 这里注意：只跟还没被“软抑制”的框算IoUious=compute_iou(boxes_copy[max_idx],boxes_copy[order[i+1:]])ifmethod=='gaussian':# 高斯衰减：IoU越大，得分乘的系数越小weights=torch.exp(-(ious*ious)/sigma)elifmethod=='linear':# 线性衰减：IoU超过阈值才衰减weights=torch.ones_like(ious)weights[ious>=0.5]=1-ious[ious>=0.5]else:raiseValueError("method must be 'gaussian' or 'linear'")# 更新得分：注意这里是逐元素乘法scores_copy[order[i+1:]]*=weights# 重新排序：得分变了，顺序也要变# 这里有个性能坑：每次循环都排序，N大时很慢# 实际工程中可以用堆排序优化order=scores_copy.argsort(descending=True)# 过滤掉得分低于阈值的框final_keep=order[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_threshold=0.5):""" boxes: [N, 4] xyxy格式 scores: [N,] diou_threshold: DIoU阈值，通常比IoU阈值设大一点，比如0.5-0.7 别这样写：直接用IoU阈值，DIoU的分布和IoU不同 """keep=[]idxs=scores.argsort(descending=True)whilelen(idxs)>0:i=idxs[0]keep.append(i)# 计算DIoU，不是IoUdious=compute_diou(boxes[i],boxes[idxs[1:]])# 抑制条件：DIoU大于阈值才抑制mask=dious<=diou_threshold idxs=idxs[1:][mask]returnkeepdefcompute_diou(box1,boxes):""" 计算box1与boxes中每个框的DIoU box1: [4] xyxy boxes: [M, 4] 这里踩过坑：坐标要归一化，否则距离计算会偏 """# 计算IoUious=compute_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_min=torch.min(box1[0],boxes[:,0])y_min=torch.min(box1[1],boxes[:,1])x_max=torch.max(box1[2],boxes[:,2])y_max=torch.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防止除零epsilon=1e-7diou=ious-d_squared/(c_squared+epsilon)returndiou

DIoU-NMS 的优势：对于包含关系（大框套小框）的情况，DIoU 比 IoU 更合理。比如一个框框住整辆车，另一个框框住车轮，IoU 可能 0.6，但中心点距离很大，DIoU 可能只有 0.2，不会被抑制。

实战对比：什么时候用哪个？

我在三个场景做过对比实验：

场景1：行人检测（密集人群）

标准 NMS：AP 72.3%
Soft-NMS（高斯，sigma=0.5）：AP 74.1%，涨了 1.8 个点
DIoU-NMS（阈值 0.6）：AP 73.5%，涨了 1.2 个点
结论：Soft-NMS 胜出，因为行人之间 IoU 高但中心点也近，DIoU 优势不明显

场景2：货架商品检测（小目标密集）

标准 NMS：AP 65.7%
Soft-NMS：AP 66.9%，涨 1.2 个点
DIoU-NMS（阈值 0.5）：AP 67.8%，涨 2.1 个点
结论：DIoU-NMS 胜出，因为商品大小不一，包含关系多

场景3：车辆检测（包含关系多）

标准 NMS：AP 78.5%
Soft-NMS：AP 79.2%，涨 0.7 个点
DIoU-NMS（阈值 0.55）：AP 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 每次循环都要重新排序，N=1000 时比标准 NMS 慢 3-5 倍。工程上可以这样优化：
- 只对得分 top-K 的框做 Soft-NMS（比如 K=200）
- 用 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）做最终过滤。这样能保留更多候选框，减少漏检。