1. 目标检测中的NMS从基础到实战目标检测任务中模型通常会输出大量重叠的预测框。想象一下你在人群中找人如果每个人都用几十个不同大小的框标记出来那画面简直没法看。这时候就需要NMS非极大值抑制来帮我们清理这些冗余框就像个智能管家只保留最合适的那个。我第一次用YOLOv5做无人机航拍图像检测时就遇到过这个问题。画面里密密麻麻的小车辆经过普通NMS处理后很多靠近的车辆直接被合并成一个了。后来发现这就是经典Hard-NMS的典型缺陷——它只关心框的重叠面积完全不管框的中心距离和长宽比。这就好比两个身高体重差不多的人站在一起硬说他们是同一个人。PyTorch里的torchvision.ops.nms就是最基础的Hard-NMS实现。它的逻辑特别直接把所有预测框按置信度从高到低排序取出当前最高分的框作为基准干掉所有和它IOU超过阈值(比如0.5)的框重复这个过程直到所有框都被处理过# PyTorch原生NMS调用示例 import torchvision keep torchvision.ops.nms(boxes, scores, iou_threshold0.5)但实际项目中你会发现当两个物体挨得很近时比如停车场里并排的车这种简单粗暴的方法很容易把正确的预测框误删。这就引出了我们需要改进的方向——不能只看重叠面积还得考虑框的几何关系。2. Hard-NMS的局限性解剖让我们做个实验在COCO数据集里找两张图一张是稀疏分布的物体另一张是密集场景。用同样的YOLOv5模型检测然后分别统计两种情况下NMS前后的预测框数量变化。你会发现密集场景下的召回率下降明显更多这就是Hard-NMS的硬伤。具体来说Hard-NMS有三大痛点中心点盲区两个框可能重叠面积很大但实际是不同物体。比如一个人抱着大纸箱人和纸箱的检测框IOU会很高。尺度不敏感一个大框和一个小框可能有相同的IOU值但明显是不同的物体。一刀切策略只要IOU超阈值就直接删除没有任何缓冲余地。在YOLOv5的源码中non_max_suppression函数最初就是直接调用PyTorch的NMS。但后来社区贡献者陆续加入了各种改进版本。我建议大家在阅读源码时重点关注general.py这个文件里面包含了各种后处理的精华代码。# YOLOv5中的NMS调用位置 def non_max_suppression( prediction, conf_thres0.25, iou_thres0.45, classesNone, agnosticFalse, multi_labelFalse ): # ... 其他预处理代码 i torchvision.ops.nms(boxes, scores, iou_thres) # 原始Hard-NMS调用 # ... 后处理代码实测在VisDrone密集目标数据集上Hard-NMS的漏检率比后面要讲的DIOU-NMS高出约15%。这个差距在无人机俯拍视角下更加明显因为这种视角下物体间距看起来更小。3. DIOU-NMS的革新之处DIOU-NMS的改进思路特别聪明——它把框之间的距离考虑进来了。DIOUDistance-IoU在原来IoU的基础上增加了一个惩罚项这个惩罚项和两个框中心点的距离成正比。公式长这样DIOU IOU - d²/c²其中d是中心点距离c是最小包围框的对角线长度。这个改进相当于给NMS加了空间感知能力两个框离得越远即使IOU相同惩罚也会越小。在YOLOv5中启用DIOU-NMS特别简单只需要在bbox_iou函数里把DIoU参数设为True# DIOU-NMS的核心实现 def bbox_iou(box1, box2, x1y1x2y2True, GIoUFalse, DIoUFalse, CIoUFalse, eps1e-9): # ... 计算IOU的基础代码 if DIoU or CIoU: c2 cw ** 2 ch ** 2 eps # 对角线长度平方 rho2 ((b2_x1 b2_x2 - b1_x1 - b1_x2) ** 2 (b2_y1 b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # 中心点距离平方 if DIoU: return iou - rho2 / c2 # DIOU公式 # ... 其他代码我在VisDrone数据集上做过对比实验同样的模型只是把NMS换成DIOU-NMS在密集行人场景下的mAP就从0.68提升到了0.73。最明显的变化是原来被合并的相邻小目标现在能正确分开了。不过要注意的是DIOU-NMS计算量会比Hard-NMS稍大因为要多计算中心点距离。在实际部署时如果硬件资源紧张可能需要做取舍。我的经验是在GPU上运行时这点开销可以忽略不计但在边缘设备上可能要测试下帧率影响。4. 代码级对比与工程实践现在我们来个实战对比看看在YOLOv5中如何切换不同的NMS方法。关键是要修改non_max_suppression函数的调用方式# 原始Hard-NMS i torchvision.ops.nms(boxes, scores, iou_thres) # 修改为DIOU-NMS iou bbox_iou(boxes[i], boxes[j], DIoUTrue) # 注意这里的DIoU参数为了更直观理解差异我整理了个对比表格指标Hard-NMSDIOU-NMS计算速度最快慢5-8%密集目标表现较差较好参数敏感性高较低部署难度简单中等在实际项目中我通常会这样做选择如果是稀疏场景如工业缺陷检测直接用Hard-NMS更高效如果是密集场景如交通监控必须用DIOU-NMS如果追求极致性能可以尝试CIOU-NMS在DIOU基础上再加长宽比考量最后分享一个调参技巧DIOU-NMS的iou_thres可以设得比Hard-NMS更高些。我一般用Hard-NMS时设0.45用DIOU-NMS会设0.6。因为DIOU的惩罚项已经能很好地区分不同物体了阈值可以放松些来提高召回。在模型导出到ONNX或TensorRT时要注意不同后处理方式的兼容性。我遇到过DIOU-NMS在TensorRT上不生效的情况后来发现是版本问题。这时候可能需要自定义插件或者回退到Hard-NMS。这些实战中的坑只有真正做过部署的人才会深有体会。