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)returndiouDIoU-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)做最终过滤。这样能保留更多候选框,减少漏检。
写在最后
NMS 这个看似简单的后处理,其实藏着很多坑。我见过有人把 Soft-NMS 的 sigma 设成 0.01,结果所有框得分都变成 0;也见过 DIoU-NMS 的阈值设成 0.3,导致大量框被误杀。调参的时候,建议先可视化几个典型场景的 IoU/DIoU 分布,心里有数再动手。
下一期我们聊聊 NMS 的进阶变体——Cluster-NMS 和 Weighted-NMS,看看怎么用聚类思想解决更复杂的重叠问题。