073、NMS 源码逐行详解：Box 筛选到按 Conf 排序到IoU 矩阵计算到抑制到保留
从一次线上事故说起
去年双十一大促，我负责的检测模型在线上疯狂输出重复框——同一个行人身上叠了七八个框，置信度还都挺高。运维同学截图发群里，配文“这模型是喝多了吗”。我第一反应是NMS写错了，结果翻出自己手撸的版本，发现确实有个低级bug：IoU计算时没做面积保护，分母为零直接崩了，但catch之后返回了0，导致所有框都被保留。从那以后，我养成了一个习惯——每次写NMS都逐行检查，尤其是边界条件。
今天就把这个“检测模型最后一道防线”的源码，从输入到输出，一行一行拆开讲。代码基于PyTorch官方实现的torchvision.ops.nms，但我会用更直白的写法来演示核心逻辑，方便你理解后再去读源码。
输入：Boxes 和 Scores
NMS的输入是两个张量：boxes形状是[N, 4]，每行是[x1, y1, x2, y2]；scores形状是[N]，对应每个框的置信度。这里有个坑——坐标必须是绝对坐标，不能是归一化后的0~1值。我之前在某个开源项目里看到有人传了归一化坐标，结果IoU算出来全是0.99，因为所有框都挤在[0,1]区间里。
defnms_custom(boxes,scores,iou_threshold=0.5):# boxes: [N, 4] (x1, y1, x2, y2)# scores: [N]# 返回保留框的索引第一步：空张量保护
ifboxes.numel()==0:returntorch.empty((0,),dtype=torch.long,device=boxes.device)别小看这行。模型推理时如果输入图像里没有目标，boxes就是空的。不处理的话，后面取索引会直接崩。我见过有人用if len(boxes) == 0，但PyTorch张量没有len，得用numel()或者size(0)。
第二步：按置信度降序排序
# 这里踩过坑：直接用scores排序，但scores可能是梯度张量# 用detach()避免梯度传播，虽然NMS不需要梯度_,order=scores.sort(descending=True)# order是降序排列的索引，比如[3, 1, 0, 2]排序是NMS的灵魂。为什么必须降序？因为我们要优先保留高置信度的框，然后用它去抑制低置信度的框。如果升序排列，你会先拿低分框去抑制高分框——逻辑上就反了。
第三步：初始化保留索引
keep=[]whileorder.numel()>0:# 每次取当前最高分的框索引i=order[0]keep.append(i)这里order会不断被截断，每次取第一个（当前最高分）加入保留列表。注意order是索引数组，不是框本身。
第四步：计算当前框与剩余框的IoU
# 如果只剩一个框，直接保留iforder.numel()==1:break# 取出当前最高分框的坐标xx1=boxes[i,0]yy1=boxes[i,1]xx2=boxes[i,2]yy2=boxes[i,3]# 取出剩余所有框的坐标（order[1:]）rest_indices=order[1:]rest_boxes=boxes[rest_indices]# 计算交集区域inter_x1=torch.max(xx1,rest_boxes[:,0])inter_y1=torch.max(yy1,rest_boxes[:,1])inter_x2=torch.min(xx2,rest_boxes[:,2])inter_y2=torch.min(yy2,rest_boxes[:,3])# 计算交集面积，这里别这样写：直接max(0, inter_w)# 因为inter_w可能是负数，表示没有交集inter_w=torch.clamp(inter_x2-inter_x1,min=0)inter_h=torch.clamp(inter_y2-inter_y1,min=0)inter_area=inter_w*inter_h这里有个细节：torch.clamp把负宽度截断为0，这样没有交集时面积就是0。有人喜欢用torch.max(inter_w, torch.tensor(0.0))，但广播机制容易出问题，clamp更干净。
# 计算当前框面积和剩余框面积area_i=(xx2-xx1)*(yy2-yy1)area_rest=(rest_boxes[:,2]-rest_boxes[:,0])*(rest_boxes[:,3]-rest_boxes[:,1])# IoU = 交集 / (并集 = 面积和 - 交集)iou=inter_area/(area_i+area_rest-inter_area)第五步：抑制低IoU框
# 保留IoU小于阈值的框索引mask=iou<=iou_threshold# 注意：mask是布尔张量，长度等于剩余框数量# 我们需要从order[1:]中选出保留的索引order=rest_indices[mask]这里order被更新为剩余框中与当前框IoU小于阈值的那些。注意rest_indices是原始索引，mask是布尔掩码，两者结合得到新的order。循环继续，直到order为空。
完整代码与边界情况
把上面拼起来就是完整的NMS。但实际工程中还有几个坑：
坐标越界：如果x2 < x1或y2 < y1，面积会是负数。虽然clamp能处理交集，但框本身的面积计算会出问题。建议在NMS之前做一次坐标校正：boxes = torch.clamp(boxes, min=0)，或者用torch.abs。
数值稳定性：当两个框完全重合时，area_i + area_rest - inter_area可能等于inter_area（因为面积相等），此时IoU=1.0，没问题。但如果框面积为零（比如x1==x2），分母为零。虽然这种情况很少见，但线上数据什么妖都有。加一个eps=1e-6是常规操作：
iou=inter_area/(area_i+area_rest-inter_area+1e-6)设备一致性：所有张量必须在同一个设备上。我见过有人从CPU加载模型，但输入是GPU张量，结果boxes和scores在不同设备上，torch.max直接报错。加一行boxes = boxes.to(scores.device)保平安。
性能优化：向量化与并行
上面是逐框迭代的写法，理解起来容易，但性能差。PyTorch官方实现用了向量化技巧——一次性计算所有框对之间的IoU矩阵，然后用循环抑制。核心思路是：
- 计算
[N, N]的IoU矩阵（上三角，因为IoU对称） - 按置信度排序后，用矩阵运算找出哪些框需要被抑制
但实际工程中，torchvision.ops.nms用的是C++实现的nms_kernel，用CUDA加速。如果你自己写Python版本，建议用torch.where和torch.any来替代显式循环，能快一个数量级。
变种：Soft-NMS 和 DIoU-NMS
标准NMS有个问题：如果两个高度重叠的目标（比如两个人挨着站），低分框会被直接干掉。Soft-NMS的做法是不直接删除，而是根据IoU降低分数：
# 不是直接删除，而是衰减分数scores[rest_indices]*=(1-iou)# 线性衰减# 或者用高斯衰减# scores[rest_indices] *= torch.exp(-iou * iou / sigma)DIoU-NMS则是在IoU基础上加入中心点距离惩罚，对遮挡场景更友好。YOLOv5和YOLOv8都支持这些变种，但默认还是标准NMS，因为简单稳定。
个人经验：调试NMS的“三板斧”
可视化中间结果：别只看最终输出。把每轮循环的
order、keep、iou打印出来，或者画图。我习惯在keep.append(i)之后，把当前框和剩余框画在同一张图上，一眼就能看出抑制逻辑对不对。阈值调参：
iou_threshold不是固定值。COCO数据集默认0.5，但实际场景要调。检测小目标时，框之间重叠少，阈值可以设高（0.7）；检测密集人群时，阈值要低（0.3~0.4）。我一般用验证集做网格搜索，步长0.05。单元测试：写一个简单的测试用例：两个完全重合的框，分数分别为0.9和0.8，阈值0.5，期望只保留高分框。再测两个完全不重合的框，期望都保留。这种测试能快速暴露bug。
最后说一句：NMS虽然只有几十行代码，但它是检测模型的“守门员”。很多模型在训练时mAP很高，上线后一塌糊涂，往往就是NMS没写好。花半小时把源码吃透，比调三天参数都值。
