1. 先说清楚为什么Iou和NMS值得单独拿出来讲做目标检测的都知道模型一轮forward出来的是几百上千个候选框而最终交付给业务方的往往只有几十个甚至几个框。中间那一步“筛选”工作靠的就是NMS非极大值抑制而NMS的排序依据、去重判据全靠IOU交并比来度量。换句话说你去任何一家做视觉相关的公司面试IOU和NMS几乎必考不是考定义而是考实现细节和边界情况。我见过不少调包调得很溜的开发者用框架自带的检测头时torchvision.ops.nms一行搞定从来没有关心过里面发生了什么。一旦换到自定义检测头、多阶段模型、或者需要部署到边缘设备的场景需要自己手写NMS或者改NMS策略时就抓瞎了。这篇文章不讲模型结构不讲训练trick就专注把IOU和NMS这两个算子讲透从数学定义、边界情况、代码实现到优化思路一次性梳理清楚。我自己最开始写这两个算子的时候也翻过车循环版本在CPU上跑慢得离谱、坐标格式没对齐导致IOU算出来全是负值、多类别NMS没有分开处理导致结果惨不忍睹。这些都是“看起来简单、做起来全是坑”的典型。希望这篇文章能帮你少走这些弯路适合刚入门目标检测的读者也适合需要手写检测后处理的工程开发者参考。2. IOU的数学本质与边界情况决定了你写出的代码是“能用”还是“好用”2.1 IOU的定义并不难难在实现时容易忽略细节IOU的全称是Intersection over Union也就是两个框的交集面积除以并集面积。公式上写出来非常简单IOU area(A ∩ B) / area(A ∪ B)其中并集面积又等于两个框面积之和减去交集面积所以计算IOU只需要算三样东西两个框各自面积、交叠区域的面积。两个框各自面积不存在什么争议就是用框的宽高相乘。真正的关键点全在“交叠区域面积”这一步。假设我们用[x1, y1, x2, y2]表示一个框其中(x1, y1)是左上角(x2, y2)是右下角那么交集区域的坐标可以这样算inter_x1 max(a_x1, b_x1) inter_y1 max(a_y1, b_y1) inter_x2 min(a_x2, b_x2) inter_y2 min(a_y2, b_y2)接着交集宽高为inter_w max(0, inter_x2 - inter_x1) inter_h max(0, inter_y2 - inter_y1) inter_area inter_w * inter_h注意这里的max(0, ...)非常关键。如果两个框完全没有重叠inter_x2 - inter_x1会是一个负数如果不加这个clip操作最后算出来的IOU会是一个负值在某些场景下还会悄悄污染后续逻辑。我之前在调试一个检测模型时发现置信度分数明明分布正常但NMS结果却偶尔出现两个完全重叠的框同时被保留查了半天最后发现问题就出在这里我当时没有对负值做clamp导致IOU算出来是负数但没有触发抑制条件。2.2 坐标格式不一致是IOU计算中最常见的隐雷不同框架、不同数据集对框的表示方式不一样表示方式含义典型场景[x1, y1, x2, y2]左上角和右下角通用检测头输出、大多数开源实现[cx, cy, w, h]中心坐标、宽高很多数据集的标注格式[x, y, w, h]左上角、宽高部分传统检测算法习惯如果直接拿cx, cy, w, h的框来做IOU计算必须先在外部统一转成[x1, y1, x2, y2]格式否则max/min逻辑完全错位。我在一个模拟项目里就经历过这种“神秘错误”训练时loss下降正常推理时mAP却很拉胯输出结果看起来“框的位置大概对了但数量明显偏多”。后来一行行打印中间过程才发现标注文件里存的是(cx, cy, w, h)而我写的后处理代码里默认是(x1, y1, x2, y2)转换关系漏了一半交叠面积全部计算错误。所以写IOU第一步不是写公式而是先明确坐标格式并且写一个显式的转换函数不要靠“记住”来隐式处理。2.3 IOU的下界是0但梯度问题并没有这么简单从数学定义上IOU的取值范围是[0, 1]。两个框完全重叠时IOU为1完全不重叠时IOU为0。但这里有一个很微妙的问题当IOU为0时交叠面积对坐标的导数是不存在的或者说梯度为0因为max(0, ...)这个操作导致局部梯度无法回传。这就是为什么后来出现了GIOU、DIOU、CIOU这一系列变体它们的主要动机之一就是解决“IOU为0时无梯度”以及“IOU相同但框的重合质量不同”这两个问题。不过在日常NMS场景里我们不需要关心梯度因为NMS本身是一个不可微的筛选操作。只有在做损失函数设计时才需要考虑IOU变体。这一点经常有人混淆我在第5节会专门展开讲。3. NMS的完整执行流程与我的第一个高效实现3.1 标准NMS到底在做什么NMS全称Non-Maximum Suppression中文叫非极大值抑制。它的核心逻辑一句话就能说清同一个物体周围通常会有多个预测框只保留置信度最高的那个然后把和它有较大重叠的框都删掉。标准的执行流程如下将所有框按置信度分数从高到低排序。取出当前分数最高的框直接保留。计算这个框与剩余所有框的IOU。删除掉IOU大于预设阈值通常0.5的框。重复第2步直到没有剩余框。整个流程非常直白但实现时有两个点需要注意。第一排序这一步必须是原始索引跟随排序否则后面保留的框索引会混乱第二删除框的操作如果顺序不对会导致“误删”和“漏删”直观表现就是检测结果中同一个物体出现多个框或者相邻的不同物体被误合并。3.2 我的第一个NMS实现以及它慢在哪先给一个朴素的PyTorch实现这种方式最符合直觉也最容易调试import torch def nms_python_loop(boxes, scores, iou_threshold0.5): boxes: (N, 4), 格式为 [x1, y1, x2, y2] scores: (N,) x1 boxes[:, 0] y1 boxes[:, 1] x2 boxes[:, 2] y2 boxes[:, 3] areas (x2 - x1) * (y2 - y1) order scores.argsort(descendingTrue).tolist() keep [] while order: idx order.pop(0) keep.append(idx) if not order: break # 当前保留框与剩余框的左上角/右下角比较 xx1 torch.clamp(x1[order], minfloat(x1[idx])) yy1 torch.clamp(y1[order], minfloat(y1[idx])) xx2 torch.clamp(x2[order], maxfloat(x2[idx])) yy2 torch.clamp(y2[order], maxfloat(y2[idx])) inter_w torch.clamp(xx2 - xx1, min0.0) inter_h torch.clamp(yy2 - yy1, min0.0) inter_area inter_w * inter_h iou inter_area / (areas[idx] areas[order] - inter_area) # 保留IOU低于阈值的框 order [i for i, v in zip(order, iou) if v iou_threshold] return keep这个实现正确性没问题理解起来也直观。但它在order.pop(0)这个地方有一个性能隐患Python列表的pop(0)是O(n)操作每次弹出一个元素都要移动后续所有元素在候选框数量大比如推理阶段一张图几百上千个框时总耗时会非常难看。3.3 用向量化重写性能提升两个数量级优化思路很简单不要真的“删除”元素而是用布尔掩码标记应该保留还是丢弃。每一步循环我们只选当前分数最高的框作为基准然后把“与基准框IOU大于阈值”的框标记为丢弃下一轮继续在剩余框里选分数最高的。一个更常见的写法如下import torch def nms_vectorized(boxes, scores, iou_threshold0.5): boxes: (N, 4) [x1, y1, x2, y2] scores: (N,) 返回保留框的索引列表 x1 boxes[:, 0] y1 boxes[:, 1] x2 boxes[:, 2] y2 boxes[:, 3] areas (x2 - x1) * (y2 - y1) order scores.argsort(descendingTrue) keep [] while order.numel() 0: if order.numel() 1: keep.append(order.item()) break idx order[0] keep.append(idx) # 计算当前最高分框与剩余所有框的IOU order_rest order[1:] xx1 torch.maximum(x1[order_rest], x1[idx]) yy1 torch.maximum(y1[order_rest], y1[idx]) xx2 torch.minimum(x2[order_rest], x2[idx]) yy2 torch.minimum(y2[order_rest], y2[idx]) inter_w torch.clamp(xx2 - xx1, min0.0) inter_h torch.clamp(yy2 - yy1, min0.0) inter_area inter_w * inter_h iou inter_area / (areas[idx] areas[order_rest] - inter_area) # 保留IOU小于等于阈值的候选框注意是候选框不是最终结果 mask iou iou_threshold order order_rest[mask] return keep这个版本把“删除”变成了掩码筛选order是Tensororder_rest[mask]是向量化索引速度比纯Python循环快很多。我在一个嵌入式推理项目里对比过候选框数量约300个时第一个版本耗时约15ms向量化版本耗时约3ms性能提升五倍。如果候选框数量上到几千个差距会拉到十几倍以上。当然工程上如果你用的是PyTorch直接调用自带的实现更省事import torchvision keep torchvision.ops.nms(boxes, scores, iou_threshold)它在底层做了一系列优化包括跨Tensor并行和更高效的数据组织方式性能通常是我们手写版本的上限。但理解它的原理依然很有价值因为很多场景下你需要修改策略比如多类别NMS、Soft-NMS这时候就只能自己动手了。4. 多类别NMS与Soft-NMS标准NMS不够用怎么办4.1 多类别NMS的“分开做”原则以及为什么不能简单交叉抑制目标检测模型的输出通常是(N, 4 1 C)其中C是类别数。如果模型同时输出多个类别NMS到底该怎么处理这里有一个常见误区直接在全部框上做一次全局NMS。不管你信不信全局NMS在大多数场景下是错的。原因很简单不同类别的物体可能合理地出现在同一位置附近比如一个“人”骑着一辆“自行车”人框和自行车框有较高IOU如果全局NMS将它们相互抑制就会丢掉其中一个正确的检测结果。正确的做法是对每个类别分别执行NMS。实现时我们需要先找到每个框所属的类别然后按类别分组处理。一个比较高效的方式是把类别信息编码进score里按(类别 置信度)的方式排序这样同一个类别的框会聚集在一起。def nms_per_class(boxes, scores, class_ids, iou_threshold0.5): boxes: (N, 4) scores: (N,) class_ids: (N,) 每个框的类别标号 keep [] for cls_id in class_ids.unique().tolist(): cls_mask (class_ids cls_id) cls_boxes boxes[cls_mask] cls_scores scores[cls_mask] keep_cls nms_vectorized(cls_boxes, cls_scores, iou_threshold) # 注意需要将局部索引映射回全局索引 global_indices cls_mask.nonzero(as_tupleFalse).squeeze(1) keep.extend(global_indices[keep_cls].tolist()) return keep这里最容易踩的坑是索引映射。cls_mask筛选后得到的索引只是子集内的局部索引如果要返回原始框上的全局索引必须通过cls_mask.nonzero()或者torch.where(cls_mask)[0]做转换。我最初实现时忽略了这一步导致返回的索引和框对不上画出来的结果张冠李戴。4.2 Soft-NMS不直接删除而是降权标准NMS的问题在于“一刀切”。当两个物体重叠程度较高时如果其中一个框的置信度略低它会直接被删除。这在拥挤场景下比如行人密集的街道效果很差。Soft-NMS的思路是不要直接删而是根据IOU大小降低该框的置信度分数。具体的做法是对与最高分框IOU超过阈值的框将它们的分数乘以一个惩罚系数。经典的Soft-NMS有两种惩罚策略线性惩罚s_i s_i * (1 - iou)。高斯惩罚s_i s_i * exp(-iou^2 / sigma)。用高斯惩罚实现时整个NMS循环里的“删除”操作被替换成“降分”操作流程变成def soft_nms(boxes, scores, sigma0.5, score_threshold0.1): x1 boxes[:, 0] y1 boxes[:, 1] x2 boxes[:, 2] y2 boxes[:, 3] areas (x2 - x1) * (y2 - y1) order scores.argsort(descendingTrue) keep [] while order.numel() 0: idx order[0] keep.append(idx) if order.numel() 1: break order_rest order[1:] xx1 torch.maximum(x1[order_rest], x1[idx]) yy1 torch.maximum(y1[order_rest], y1[idx]) xx2 torch.minimum(x2[order_rest], x2[idx]) yy2 torch.minimum(y2[order_rest], y2[idx]) inter_w torch.clamp(xx2 - xx1, min0.0) inter_h torch.clamp(yy2 - yy1, min0.0) inter_area inter_w * inter_h iou inter_area / (areas[idx] areas[order_rest] - inter_area) # 高斯惩罚分数乘以 exp(-iou^2 / sigma) weights torch.exp(-(iou * iou) / sigma) scores[order_rest] * weights # 丢弃分数过低的框重新排序 order order_rest[scores[order_rest] score_threshold] order order[scores[order].argsort(descendingTrue)] return keepSoft-NMS在密集场景下通常能带来稳定的mAP提升代价是计算量比标准NMS略高因为同一个框可能被反复加权。实际部署时需要根据业务对延迟的敏感度来权衡。4.3 NMS阈值的玄学0.5、0.6、0.3分别适合什么场景NMS的iou_threshold是一个超参数它直接控制“多大重叠算重复”。我整理的参考经验如下阈值行为特点适用场景0.3抑制力度大重叠稍多就删稀疏目标场景对误检容忍度低0.5最常用均衡通用检测任务默认值0.7抑制力度弱保留更多框密集小目标场景重叠目标多的场景如果做的是密集人群计数、细胞检测这类任务阈值可以调高到0.7否则同一目标很容易因为相互重叠被误删最终的框数量会明显偏少。如果做的是工业质检这类稀疏场景误检率会影响准确率可以适当调低。需要特别注意的是阈值调参没有“万能最优解”不同模型、不同数据分布下的最优点差异很大。最靠谱的做法是在验证集上做一个小网格搜索把NMS阈值当成一个超参数参与调优。5. 从IOU到GIOU、DIOU、CIOUNMS之外的那些“亲戚”5.1 IOU作为损失函数的两大痛点前面提到计算IOU时用了clamp(0)导致梯度为0这只是表面现象。更深层的痛点是当两个框完全不重叠时IOU恒为0模型无法通过梯度感知“相距多远”。即使IOU相同不同情况的对齐质量完全不同但IOU无法区分。举一个非常典型的例子假设有两个框它们的IOU都等于0.5但一种是“中心对齐、大小不同”另一种是“中心偏了但大小匹配”。从目标检测的质量来说前者通常比后者好因为中心点偏差小IoU度量却完全一样。这就是GIOU等变体出现的动机。5.2 GIOU引入最小外接框来提供“距离感”GIOUGeneralized IoU的做法是在两个框的交并比基础上引入包含它们的最小外接矩形C并计算C - (A ∪ B)与C的比值作为一个惩罚项GIOU IOU - (C - (A ∪ B)) / C这个式子的直观含义是两个框越分散最小外接矩形C越大惩罚项越大即使IOU为0GIOU也有一个可用的值趋近于-1梯度就不再为0了。它的取值范围是[-1, 1]。5.3 DIOU与CIOU把中心距离和宽高比纳入考量DIOUDistance-IoU在GIOU的基础上更进一步直接惩罚两个框中心点的归一化距离DIOU IOU - (中心距离^2) / (最小外接框对角线距离^2)当两个框中心点重合时DIOU退化为IOU。而CIOUComplete-IoU在DIOU的基础上又增加宽高比一致性惩罚项是目前目标检测损失函数里使用最广泛的。指标惩罚项内容主要应用场景GIOU最小外接框面积差弥补IOU为零时的梯度缺失DIOU中心点距离加速框的收敛速度CIOU中心点距离 宽高比更精细的边界框回归损失需要注意的是GIOU、DIOU、CIOU一般是用在训练损失里而不是用来替代NMS中的IOU计算。NMS关心的只是“两个框重叠程度”标准IOU已经足够且高效。有人可能误以为用CIOU替换NMS里的IOU效果会更好实际上NMS阶段没有梯度参与这些改进指标的理论优势根本没有发挥空间反而白白增加计算开销。5.4 我踩过的坑在损失函数里直接用坐标差导致收敛慢曾经在一个模拟项目里我偷懒直接用坐标差的L2损失做框回归结果模型训练到中期就卡住了loss降不动检测出来的框中心点偏、尺寸也不准。后来换成CIOU损失问题立刻缓解。原因就是L2损失对不同尺度的框不敏感而CIOU对中心点距离和宽高比的惩罚更符合目标检测的实际需求。这里给一个PyTorch风格的CIOU计算示例方便对照参考def ciou_loss(pred_boxes, target_boxes): pred_boxes, target_boxes: (N, 4) [cx, cy, w, h] pred_cx, pred_cy, pred_w, pred_h pred_boxes.unbind(dim1) target_cx, target_cy, target_w, target_h target_boxes.unbind(dim1) # 计算IOU pred_x1 pred_cx - pred_w / 2 pred_y1 pred_cy - pred_h / 2 pred_x2 pred_cx pred_w / 2 pred_y2 pred_cy pred_h / 2 target_x1 target_cx - target_w / 2 target_y1 target_cy - target_h / 2 target_x2 target_cx target_w / 2 target_y2 target_cy target_h / 2 inter_x1 torch.maximum(pred_x1, target_x1) inter_y1 torch.maximum(pred_y1, target_y1) inter_x2 torch.minimum(pred_x2, target_x2) inter_y2 torch.minimum(pred_y2, target_y2) inter_area torch.clamp(inter_x2 - inter_x1, min0) * torch.clamp(inter_y2 - inter_y1, min0) pred_area pred_w * pred_h target_area target_w * target_h union_area pred_area target_area - inter_area iou inter_area / (union_area 1e-6) # 中心点距离惩罚 center_dist (pred_cx - target_cx) ** 2 (pred_cy - target_cy) ** 2 # 最小外接框对角线距离 close_x1 torch.minimum(pred_x1, target_x1) close_y1 torch.minimum(pred_y1, target_y1) close_x2 torch.maximum(pred_x2, target_x2) close_y2 torch.maximum(pred_y2, target_y2) close_diag (close_x2 - close_x1) ** 2 (close_y2 - close_y1) ** 2 1e-6 diou iou - center_dist / close_diag # 宽高比一致性惩罚 v (4 / (3.1415926 ** 2)) * ((torch.arctan(target_w / (target_h 1e-6)) - torch.arctan(pred_w / (pred_h 1e-6))) ** 2) alpha v / (1 - iou v 1e-6) ciou diou - alpha * v return 1 - ciou这里的1e-6是防止除零。另外要小心atan2和arctan的区别我一开始用错函数导致宽高比惩罚在几个特殊角度上跳变训练过程不稳定调整后才恢复正常。6. 实战中的性能对比与选型建议写给准备上手的人6.1 用一张表对比标准NMS、Soft-NMS和自实现的取舍方案实现难度推理耗时约300框精度收益适用场景torchvision自带的NMS极低1ms基准绝大多数通用检测任务手写向量化NMS中3ms左右与基准持平需要定制逻辑、无法依赖框架时Soft-NMS中高5ms左右密集场景有提升行人、车辆拥挤场景多类别NMS中视类别数而定避免跨类别误删多类别检测模型可以看出PyTorch自带的NMS实现基本是性能上限除非有定制化需求比如修改惩罚函数、添加类别约束、或做一些逻辑上的改动否则没必要手写并替换。6.2 我在一个低算力设备上的部署经验在一个边缘设备的部署项目里我遇到了一个很有意思的问题模型本身已经压缩过单帧推理约25ms但nan的IOU计算和NMS在CPU上跑了接近12ms占比接近三分之一有点无法接受。后来排查发现原因有两个一是当时用了纯Python嵌套循环的NMS实现二是在NMS之前对每个框都计算了IOU矩阵而实际上NMS只需要逐步计算最大框与剩余框的IOU没必要算全量矩阵。优化方案是把NMS改为直接使用框架自带的nms算子很多推理框架都有对应的融合算子。如果必须自实现也应该避免生成完整的(N, N)IOI矩阵而是每轮只生成(N-1,)的一维张量。多类别时不要按类别循环调用NMS而是一次把所有框送入NMS再在返回索引时结合类别信息做过滤。这三步做完之后NMS耗时从12ms降到了3ms左右整体推理速度提升了近15%。6.3 如果输入坐标是归一化的要不要做特殊处理很多数据预处理会把坐标归一化到[0, 1]此时IOU计算本身不受影响因为公式里只用到坐标差值缩放不会改变IOU的值。但有一点要注意如果框的坐标在某些样本中出现了负值比如增强时越界裁剪在计算宽高之前一定要clamp(0)否则面积会算出错误结果。另外归一化坐标在可视化时容易忽略反归一化而画出一个贴近边缘的框这个问题不影响IOU/NMS逻辑但在调试时非常迷惑人建议调试时统一用原始像素坐标。6.4 我最后想分享的一个小技巧在调试NMS逻辑时建议先在单类别、少量框的场景里逐行print每一步的order、keep、iou结果手动核对一遍流程再放到真实数据上跑。我在写Soft-NMS的初始版本时就因为排序后没有重新更新order索引导致程序稳定地返回错框靠print定位了两小时才发现是“降权后忘记重新排序”的问题。这个错误写上几行就能避免但确实很隐蔽希望你不要踩一遍。最后再补充一点IOU和NMS虽然看起来基础但它们是检测头之后的第一道关卡任何偏差都会被下游的评估指标放大。调模型时如果mAP一直上不去不妨回头检查一下这两段计算逻辑也许真正的瓶颈根本不在模型结构里。