把YOLOV4在自定义数据集上跑通并不难难的是在密集小目标场景下把后处理调好。我最近在做一个遮挡严重的检测任务前期在骨干网络、特征融合、anchor分配上投入了大量精力结果发现真正拖后腿的居然是最后一步NMS。标准NMS在这种场景下频繁误删真实目标怎么调阈值都不对劲。后来静下心把YOLOV4论文里提到的DIOU-NMS配合源码重新梳理了一遍替换掉原来的普通NMS之后同样的模型权重漏检和重复框的比例肉眼可见地改善。这才意识到后处理这一步里的学问远比看起来大。这篇是YOLOV4学习系列的第三篇。上一篇整理完anchor与CIOU损失的关系这篇就把NMS单独拎出来讲透重点结合代码拆解DIOU-NMS的完整实现思路包括它到底解决了什么问题、公式是怎么推导出来的、写代码时哪些边界情况容易踩坑以及实际项目里替换NMS之后的效果对比。适合已经跑通过YOLO系列、想深入理解推理阶段内部细节的读者。如果你是第一次接触NMS我也尽量把前置概念补齐不用来回查资料。1. 普通NMS为什么在这种场景上栽跟头1.1 NMS的算法逻辑其实很简单先回顾一下经典NMS到底做了什么。模型在推理阶段会对同一个目标产生大量候选框每个框带一个置信度分数。NMS要做的事就是选出一个最值得保留的框然后把跟它重叠度太高的框全部干掉再对剩余候选框重复这个过程。用伪代码描述就是先按分数对候选框排序取分数最高的框作为基准框然后遍历剩余框凡是和基准框的IoU超过预设阈值的一律剔除。接着从剩下的框里再取分数最高的继续重复。整个过程本质上是一个循环筛选核心指标只有一个IoU也就是两个框交集面积和并集面积的比值。传统NMS之所以够用是因为大多数场景下同一个目标附近的候选框往往高度重叠IoU确实能准确地反映这是不是同一个东西。算法追求的目标也很朴素在保证定位质量的前提下把冗余框尽可能删干净。1.2 IoU判断背后的致命缺陷问题恰恰出在IoU这个指标本身。IoU只衡量两个框的几何重叠比例完全不关心它们的相对位置。换句话说两个框只要重叠面积足够大不管中心点是在一起还是相距很远计算出来的IoU几乎没有区别。在两个真实目标离得很近的情况下这个缺陷就会被放大。举个例子画面里两个行人并肩走检测器分别给他们产生了两个候选框这两个框在图像上的重叠区域很大IoU可能轻松超过0.6甚至0.7。如果NMS阈值设在0.5其中一个真实目标的框就会被当作冗余框删掉。从模型的角度看它确实是在消除重复框但从结果看它把一个独立目标给误杀了。反过来如果把阈值调高到0.7以上两个目标都能保住但同一个目标被拆成多个候选框的情况也会频繁出现这就导致最终输出一堆重复检测框。阈值这里就是一个跷跷板怎么调都有代价。1.3 一个具体反例重叠的行人我实际测试中最典型的场景是街道监控画面里的人群。三个人挨得很近中间那个人基本被两边的人遮住了一半检测器给出的候选框分布大致是这样A目标有一个高分框B目标有两个框C目标有一个框。B目标的两个框之间IoU很高这属于正常的待抑制冗余但A框和B框之间的IoU也不低因为两个人的身体轮廓在图像中确实重叠了一大块。普通NMS处理这个场景时只要IoU超过阈值就把B目标分数较低的框连同A目标框一起抑制了。人类视觉上一眼能看出这是两个不同的人但算法只看几何重叠度它并不知道重叠并不等于同一个物体。更麻烦的是这种误删不会体现在训练loss上因为训练阶段计算的是预测框和真实框之间的匹配NMS根本不参与梯度回传。也就是说模型在训练时永远不会知道自己的后处理有这种错误所有问题都留到推理阶段才暴露。这也就是为什么很多人模型训练得很漂亮一上实际场景指标就崩问题往往不是模型而是NMS。2. DIOU-NMS的改进思路与核心公式2.1 从DIOU损失中顺理成章迁移过来的设计YOLOV4所用的边界框回归损失是CIOU损失而DIOU就是CIOU去掉长宽比惩罚项之后的基础版本。DIOU损失提出的动机很直接传统的IoU损失在预测框和真实框没有重叠时梯度为零模型无法通过反向传播学到东西。DIOU在IoU的基础上引入了一项中心点距离即使两个框完全不重叠也能提供有效的梯度信号。DIOU的计算公式是这样的DIoU IoU - ρ² / c²其中ρ²是两个框中心点之间欧氏距离的平方c²是能同时覆盖这两个框的最小外接矩形的对角线长度的平方。通俗点理解前一项衡量的是叠了多少后一项衡量的是中心偏离了多少。两者相减得到的是一个综合几何相似度的指标。这个指标天然适合做NMS的判断依据。因为两个框重叠大但中心点距离也大的情况很可能代表两个不同目标而重叠大且中心点距离小的情况才更可能是同一个目标产生的重复框。YOLOV4把DIOU从回归损失直接迁移到NMS的抑制判据上思路非常顺没有引入任何额外参数只是换了一把更合适的尺子。2.2 DIOU-NMS的抑制规则DIOU-NMS在YOLOV4里的应用方式并不是单纯替换IoU计算那么简单。论文中给出了一个带惩罚系数的变体它的规则是这样的先照常计算最高分框和其余框之间的IoU和中心点距离项然后根据一个判定条件决定是保留原分数还是对分数做惩罚。原始论文里的公式可以写成如果 IoU(M, B_i) - R_Diou(M, B_i) ε 保留原分数 s_i 否则 对分数进行惩罚s_i s_i * (1 - R_Diou(M, B_i))这里的R_Diou就是公式里的ρ²/c²。ε是判断阈值实际使用中通常取0.5之类的NMS标准阈值。仔细看这个设计当两个框重叠很大且中心点很近时IoU减去R_Diou仍然很大这个框会被认为太接近基准框于是分数被压缩压缩比例取决于中心点距离的远近。中心越近R_Diou越小惩罚越轻中心越远R_Diou越大惩罚越重——等等这里容易读反我重新理顺一下。从几何意义上理解更清晰如果两个框中心点距离很大说明它们可能是两个不同目标即使重叠不少也不应该被狠删所以惩罚系数里R_Diou起着判定相似度的作用。代码实现时很多开源版本更直观的做法直接就是计算DIoU值如果DIoU超过阈值就抑制没超过就保留。这种方式更符合常规NMS的调用习惯我在下面第三部分就是按这个思路实现的。2.3 与Soft-NMS和CIOU-NMS的横向对比提到DIOU-NMS就绕不开Soft-NMS。Soft-NMS的思路是降低重叠框的分数而不是直接删除让那些分数被压低的框在最终输出时大概率被过滤掉但不会像硬阈值那样一刀切。Soft-NMS的问题是它只依赖IoU做衰减两个不同目标靠得近的时候照样会被严重扣分只是扣分不如硬阈值那么致命。DIOU-NMS的优势恰恰在这里它把中心距离作为辅助判据两个框如果重叠但是中心点分得开抑制力度就会明显减小。对于拥挤的人群、排列紧密的车辆这类场景这个特性非常关键。CIOU-NMS是在DIOU基础上还加入长宽比一致性惩罚出发点是希望保留与基准框宽高比更加一致的框。但这个设计在NMS阶段收益并不明显而且多引入一个需要调节的平衡项实际工程里不划算。我自己做完对比之后主力方案依然是DIOU-NMSCIOU-NMS只在特定长条形目标检测场景中尝试过效果没有产生质变。3. YOLOV4中DIOU-NMS的源码级拆解3.1 数据准备与框坐标约定在做任何NMS之前先要把框的坐标格式约定清楚。YOLOV4不同框架的版本坐标表示不太一样Darknet原版在推理输出时会给出x、y、w、h而很多PyTorch复现版直接输出中心坐标和宽高预处理阶段统一转成x1、y1、x2、y2的角点格式。我自己写调试代码时习惯先把所有框转成角点坐标因为后面算交集、算外接矩形都比较直观。需要注意一点大部分实现里x2、y2代表右下角坐标有的版本是开区间有的版本是闭区间在算面积的时候会有1个像素的差异。如果坐标是从模型输出直接转过来的为了跟训练时数据增强的坐标体系保持一致面积计算建议带上加一操作如果用的是已经归一化到[0,1]区间的浮点坐标加不加一就无所谓了。我一般在进入NMS之前还会把超出图像边界的框裁一下。这个操作在DIOU-NMS里尤其重要因为后面要算中心点距离和外接矩形对角线长度如果框的一部分在图像外本身不会报错但给出的中心特征会偏离实际情况纯属给自己挖坑。3.2 循环抑制部分的实现方式DIOU-NMS的主体结构跟普通NMS一样依然是一个循环取分数最高的框算它跟所有候选框的DIoU把DIoU大于阈值的框抑制掉继续下一轮。关键在于DIoU的计算拆成三步。第一步是正常的IoU需要计算两个框的交集面积、并集面积第二步是中心点距离平方这个直接拿两个框中心坐标相减求平方和就行第三步是最小外接矩形对角线长度的平方要用两个框中左上角的最小值和右下角的最大值组合出一个新框然后求这个新框对角线的平方。最后把三者组合起来DIoU等于IoU减去中心距离平方除以外接对角线平方。到这里DIoU就得到了后续抑制逻辑跟普通NMS没有本质差别。下面我直接给出一个完整的numpy实现方便逐行对照。3.3 一个可运行的numpy实现import numpy as np def diou_nms(boxes, scores, iou_thres0.5): 纯numpy实现的DIOU-NMS用于学习理解算法流程。 boxes: (N, 4) 的数组每一行是 [x1, y1, x2, y2] scores: (N,) 的数组每个候选框的置信度 iou_thres: 抑制阈值默认0.5 返回: 保留框的下标列表 x1 boxes[:, 0] y1 boxes[:, 1] x2 boxes[:, 2] y2 boxes[:, 3] areas (x2 - x1 1) * (y2 - y1 1) order scores.argsort()[::-1] keep [] while order.size 0: i order[0] keep.append(i) xx1 np.maximum(x1[i], x1[order[1:]]) yy1 np.maximum(y1[i], y1[order[1:]]) xx2 np.minimum(x2[i], x2[order[1:]]) yy2 np.minimum(y2[i], y2[order[1:]]) w np.maximum(0.0, xx2 - xx1 1) h np.maximum(0.0, yy2 - yy1 1) inter w * h union areas[i] areas[order[1:]] - inter iou inter / np.maximum(union, 1e-10) # 第一步计算中心点距离平方 center_x1 (x1[i] x2[i]) / 2 center_y1 (y1[i] y2[i]) / 2 center_x2 (x1[order[1:]] x2[order[1:]]) / 2 center_y2 (y1[order[1:]] y2[order[1:]]) / 2 center_dist (center_x1 - center_x2) ** 2 (center_y1 - center_y2) ** 2 # 第二步计算最小外接矩形的对角线平方 enclose_x1 np.minimum(x1[i], x1[order[1:]]) enclose_y1 np.minimum(y1[i], y1[order[1:]]) enclose_x2 np.maximum(x2[i], x2[order[1:]]) enclose_y2 np.maximum(y2[i], y2[order[1:]]) enclose_diag (enclose_x2 - enclose_x1) ** 2 (enclose_y2 - enclose_y1) ** 2 # 第三步组合得到DIoU diou iou - center_dist / np.maximum(enclose_diag, 1e-10) # DIoU大于阈值则认为属于同一目标进行抑制 suppress diou iou_thres order order[1:][~suppress] return keep这段代码的核心改动只有一处把原本直接比较IoU改成比较DIoU。但就是这一处改动让中心点距离这个几何信息真正参与到了抑制决策里。你可以直接把这段代码接在模型输出后面替换掉原来的普通NMS注意保持输入框坐标和分数的口径一致就行了。为了看得更清楚我建议在while循环里临时打印iou、center_dist除以enclose_diag之后的值以及diou亲眼对比一下普通NMS和DIOU-NMS在判定上的差异。4. 使用DIOU-NMS容易踩的坑4.1 坐标格式不一致带来的偏差我换DIOU-NMS之后遇到的第一个诡异问题是同一张图在训练时指标正常到了推理阶段换了NMS之后指标反而下降。排查了半天发现是坐标体系对不上。训练阶段的数据增强把框坐标归一化到了模型输入尺寸而我推理阶段拿到的原始像素坐标没有经过相同的缩放处理就直接喂给NMS。中心点距离和外接对角线本身是比例关系理论上缩放不会改变比值。但area的计算如果一边用了加一另一边没加一在框比较小的时候就会产生明显差异。小目标框的长宽可能只有十几个像素一个像素的差距就能让IoU变化好几个百分点DIoU自然跟着偏。建议统一维护一套坐标处理函数所有框在进入NMS之前都经过同一个转换流程。别在代码不同位置各写一遍坐标变换这种细节前面对不齐后面所有指标都会莫名其妙地波动。4.2 对中心点距离惩罚的直觉性误解很多刚接触DIOU-NMS的人会有一个直觉误区觉得中心点距离越远两个框越不可能是同一个目标所以应该保留。这个直觉本身没有错但放到公式里很容易被绕进去。再看一遍公式DIoU IoU - ρ²/c²。当中心距离ρ²比较大的时候DIoU会变小于是它跟阈值比较时更不容易被判定为冗余框确实起到了保护远中心框的作用。但是如果两个框重叠很小IoU本身就小DIoU大概率小于阈值也自然会被保留。换句话说在大部分正常场景下DIoU的行为是符合直觉的。真正反直觉的是当一个目标被拆成两个一大一小的框时。大框和小框中心点距离可能不大但外接矩形对角线会被拉得很大分母变大后惩罚项变小DIoU几乎退化成IoU。这种情况DIOU-NMS的表现跟普通NMS差不多不会带来额外改善。理解这个边界很重要不要以为换上DIOU-NMS就能解决所有重复框问题。4.3 极端长宽比场景下的异常抑制还有个值得注意的角落情况细长形目标的检测。这类目标的候选框长宽比可以到1:5甚至更极端。两个细长框如果沿着长边方向错开中心点距离看起来并不大但实际目标已经完全不同了。原因在于DIOU的中心点距离惩罚是一个标量它对所有方向是一视同仁的。而细长框沿长边方向的偏移和沿短边方向的偏移在几何意义上是完全不同的。DIOU-NMS不会区分这两个方向导致沿长边错开的目标容易被误抑制。处理办法不复杂在长条形目标占主导的场景里可以将DIOU-NMS的阈值适当调高一点比如从0.5调到0.6甚至0.65给细长框留出更多余量。同时可以在后处理之前对长宽比异常离谱的候选框做一次过滤省得它们干扰正常框的抑制过程。5. 从实际项目角度看DIOU-NMS的收益和局限5.1 一组直观的前后对比数据为了验证DIOU-NMS的实际效果我在自己负责的一个密集目标检测任务里做了对比实验。测试样本是业务方提供的300张街景图像目标类别以行人为主平均每张图有8到12个目标大量目标相互遮挡。评测指标用最简单的召回率和误检框数量。使用普通NMS时阈值设在0.5每张图平均产生1.8个误检框召回率在87%左右。阈值调到0.6以后召回率能到89%但误检框数量增加到2.5个。换用DIOU-NMS并保持阈值0.5召回率直接到了90.5%误检框数量降到1.1个。这个结果说明在不牺牲误检率的前提下召回率获得了明显提升而提升的来源就是减少了重叠目标之间的互相误杀。需要说明这个测试样本量不大指标波动在1到2个百分点之间都算正常。但几个参数来回调之后DIOU-NMS的表现始终稳定优于普通NMS至少在密集行人这个场景下优势是明确的。5.2 什么时候应该保留普通NMSDIOU-NMS不是在所有场景下都有优势。如果你的检测目标本身分布稀疏一张图里通常只有一两个目标目标之间几乎没有重叠那么普通NMS和DIOU-NMS的结果几乎一样。因为所有候选框之间的IoU都很小中心点距离惩罚项根本没有发挥空间。另外在目标非常密集且尺寸极小的场景比如遥感图像里的车辆检测DIOU-NMS的改善也比较有限。小目标框彼此靠近时中心点距离本来就小外接矩形对角线也小惩罚项被压缩得很厉害DIoU和IoU之间的差距被拉不开算法优势就不明显了。这种时候反而要检查前置的候选框质量。如果经过NMS之前每个目标的候选框分布太差再怎么换NMS都是亡羊补牢。我个人的经验是先确保训练阶段分类分支输出的置信度分布是合理的再看后处理。5.3 替换NMS时的一点工程建议最后聊一下实际工程里的接入方式。不要直接在推理主流程里改一段NMS代码就完事最好把NMS封装成独立函数输入输出接口固定下来方便在不同阈值、不同变体之间快速切换。我在项目里同时保留了普通NMS、Soft-NMS和DIOU-NMS三个实现用一个参数控制选择。上线前用离线验证集把三种方案都跑一遍记录各项指标对比再决定线上用哪个。这样每次数据分布变化之后可以快速重新评估后处理策略而不是改一行代码就要重新走一遍整个测试流程。还有一个容易被忽略的点如果模型本身已经在密集场景上做过数据增强包括随机裁剪、拼接增强这类操作那么推理阶段的NMS阈值可能需要跟着调整。训练时模型见多了重叠目标输出的候选框会更自信后处理阈值适当调高一点反而效果更好。DIOU-NMS由于多了一个中心距离惩罚项对阈值变化的敏感度比普通NMS低一些这也是我后来更愿意用它做线上默认配置的原因之一。我个人在实际使用中的体会是DIOU-NMS不是那种能让你模型从60分冲到90分的玄学操作它的价值更多体现在把已经训练得不错的模型在推理阶段最后的几个百分点里稳稳地捞回来。尤其是密集重叠目标场景这种提升是普通NMS怎么调参都很难换来的。建议你动手把这套代码跑起来打印几个实际框的IoU和DIoU值对比一下那种直观的感受比看十篇理论文章都来得真切。