小目标检测老漏检把 YOLOv11 的损失函数换成 NWD 试试做目标检测的同行应该都有这种体会模型在常规尺寸目标上表现挺好一遇到航拍图、医学影像、工业质检里的小目标mAP 掉得惨不忍睹。这不是模型容量不够往往是损失函数在“拖后腿”。我之前在红外小目标数据集上踩了不少坑最后把 YOLOv11 默认的 CIoU 损失换成 NWDNormalized Wasserstein Distance损失小目标检测精度提升非常明显。这篇博客就把我完整的改进思路、代码实现和实验对比写出来给被小目标检测困扰的朋友一个可直接参考的方案。需要提前说明的是这套改进思路适合正在用 YOLOv11 做目标检测、尤其是数据集中小目标占比较高的同学。如果你只是跑跑通用检测任务目标尺寸没有明显偏向小目标那默认配置就够了不必折腾。1. 项目概述1.1 问题拆解小目标检测难在哪小目标检测难这是一个老生常谈但值得反复琢磨的问题。从实际数据来看小目标通常指像素尺寸小于等于 32x32 的目标或者是其面积占图像总面积比例极低的目标。这类目标在检测上有几个固有难点第一是特征信息太少。小目标在特征金字塔的不同层上经过下采样后可能只剩下几个甚至一个像素点能提取到的语义信息极其有限。第二是锚框匹配困难小目标的 IoU 对偏移极其敏感。这里解释一下IoU 计算的是预测框和真实框的交并比当目标框本身很小时哪怕预测框只偏移了一两个像素IoU 也会急剧下降。举个例子一个 10x10 的真实目标框如果预测框偏移了 2 个像素IoU 可能直接从 0.9 掉到 0.5 以下。而一个 100x100 的大目标框偏移 2 个像素对 IoU 的影响微乎其微。这就导致小目标在训练过程中很难被分配为正样本模型自然学不好。1.2 方案选型为什么从损失函数入手面对小目标检测精度低的问题常见的改进路径有几条改进特征金字塔结构、添加注意力机制、优化锚框分配策略、更换数据增强方式以及改损失函数。我之前试过在 YOLOv11 中添加 CBAM 注意力机制也试过用 BiFPN 替换原有 Neck 结构效果有提升但带来的计算量开销不小而且对训练调参比较敏感。后来阅读了不少文献发现 NWDNormalized Wasserstein Distance损失函数在小目标检测上有很好的理论基础专门针对 IoU 对小目标敏感的问题。它的核心思路是用 Wasserstein 距离来度量预测框和真实框之间的分布距离并且对尺度不敏感。换句话说它天然适配小目标场景。从改动成本来看在 YOLOv11 中把边界框损失从 IoU 系列换成 NWD 损失只需要改动损失函数计算部分不需要改动网络结构也不影响推理速度属于性价比很高的一种改进。这也是我最终选择这个方案的原因。2. 原理拆解2.1 YOLOv11 的损失函数构成要替换损失函数得先弄清楚 YOLOv11 原来的损失是怎么算的。YOLOv11 延续了 YOLOv8 的整体架构其检测头损失主要由三部分组成边界框回归损失box_loss、分类损失cls_loss和 DFL 损失Distribution Focal Loss。其中边界框损失采用的是 CIoU 损失分类损失采用 BCE 损失DFL 损失的作用是让模型预测的目标框边界分布更加集中主要优化边界框四条边的位置分布。这三部分损失按照一定权重相加就得到了最终的总损失。具体来说YOLOv11 最终的损失函数可以表示成这样Loss λ1 * L_cls λ2 * L_box λ3 * L_dfl在 YOLOv11 默认配置中这三个权重系数由模型配置文件和训练超参数共同决定。通常分类损失和 DFL 损失保持默认而 box_loss 的优化函数可以替换。CIoU 损失是在 IoU 基础上加入中心点距离和长宽比惩罚项公式展开后可用以下伪代码示意# 伪代码示意不代表ultralytics源码 def ciou_loss(pred_box, target_box): iou compute_iou(pred_box, target_box) center_distance compute_center_distance(pred_box, target_box) aspect_ratio_penalty compute_aspect_ratio_penalty(pred_box, target_box) loss 1 - iou center_distance aspect_ratio_penalty return loss这个损失函数在通用目标检测上表现很好收敛快、定位准。但你仔细观察就会发现它的第一项仍然包含 IoU 计算只要目标框较小IoU 对偏移的敏感性依然存在。2.2 NWD 的数学原理NWD 的全称是 Normalized Wasserstein Distance即归一化沃瑟斯坦距离。它的思想是把边界框建模为二维高斯分布然后计算这两个高斯分布之间的 Wasserstein 距离。Wasserstein 距离本身是度量两个概率分布之间差异的一种方式相比 IoU 这种几何重叠度量的优势在于它对两个分布是否重叠不那么敏感即使两个框完全不重叠也能计算出一个有意义的距离值。具体做法是把一个边界框视为以框中心为均值、以框的宽高为方差的高斯分布。给定两个边界框 A 和 B分别建模为高斯分布它们之间的二阶 Wasserstein 距离存在解析解W2(A, B) ||(cx_a, cy_a, w_a/2, h_a/2) - (cx_b, cy_b, w_b/2, h_b/2)||^2这个公式的结果是一个距离值衡量两个高斯分布的差异。但直接使用会有尺度问题因为距离值会随着目标尺寸增大而增大不利于不同尺度目标的统一训练。于是论文提出了归一化处理使用指数函数将距离映射到 0 到 1 之间NWD(A, B) exp(-sqrt(W2(A, B)) / C)其中 C 是一个与数据集相关的常数通常设置为数据集中目标平均尺寸的某个比例。经过这种归一化后NWD 的值范围是 0 到 1值越大表示两个框越接近。2.3 为什么 NWD 对小目标更友好从原理上看NWD 之所以对小目标检测友好核心在于它消除了 IoU 对目标尺寸的敏感性问题。为了帮助理解我常打一个比方IoU 就像用两把尺子去量两条线段的重合比例如果线段本身很短哪怕两把尺子只是稍微错开了一点重合比例也会急剧下降。而 NWD 更像是比较两条线段中心点之间的距离和各自长度的综合差异它考虑的是两个分布的相似程度而不是几何重叠比例。正因为如此在训练小目标时NWD 不会因为几个像素的偏移就产生极大的损失值波动梯度更新更平滑。同时在正负样本分配阶段NWD 也能更稳定地反映预测框和真实框的接近程度减少小目标漏检。从实验结果看许多研究都验证了在小目标检测数据集上NWD 相比 IoU 系列损失有明显的 mAP 提升。而且 NWD 与 CIoU 可以配合使用先用 NWD 辅助训练再用 CIoU 精细化优化两者并不冲突。3. 代码落地3.1 ultralytics 检测头源码结构实际动手改代码之前要先看懂 YOLOv11 的损失计算在 ultralytics 中是如何组织的。如果你用的是 ultralytics 官方库检测头损失的主要代码在ultralytics/utils/loss.py文件中核心是DetectionLoss类。这个类有以下几个关键的组成部分bce是分类损失的计算模块使用 BCEWithLogitsLosshierarchical_cls是 YOLOv11 新增的分层分类结构DFL是分布聚焦损失的实现box_loss方法负责计算边界框回归损失。在DetectionLoss类的__call__方法中会先通过self.preprocess处理预测结果然后计算分类损失和 DFL 损失接着调用self.box_loss计算边界框损失最后按权重合并输出。具体的代码片段大致如下不同版本可能有细微差异loss torch.zeros(3, deviceself.device) # box, cls, dfl # 分类损失 loss[1] self.bce(pred_scores, target_scores).sum() / target_scores.shape[0] # DFL损失 loss[2] self.dfl(pred_dfl, target_dfl).sum() / target_scores.shape[0] # box损失默认使用CIoU loss[0] self.box_loss(pred_boxes, target_boxes, target_scores.sum(-1))所以如果要把 CIoU 替换成 NWD 损失核心就是把box_loss方法内部的 IoU 计算方式替换为 NWD 计算方式。3.2 在 YOLOv11 中添加 NWD 损失函数现在开始动手实现。第一步需要把 NWD 的计算函数添加到loss.py中。这里我直接贴出我改造后的代码import torch import torch.nn as nn def wasserstein_loss(pred_boxes, target_boxes, eps1e-7, constant12.8): 计算归一化Wasserstein距离损失 参数: pred_boxes: 预测框格式为 (N, 4)每行为 [cx, cy, w, h] target_boxes: 真实框格式为 (N, 4)每行为 [cx, cy, w, h] constant: 归一化常数论文中建议设置为数据集中目标平均尺寸相关值 返回: nwd_loss: 归一化Wasserstein距离损失的均值 # 计算中心点距离的平方 center_dist (pred_boxes[:, :2] - target_boxes[:, :2]) ** 2 # 计算宽高差的平方 wh_dist (pred_boxes[:, 2:] - target_boxes[:, 2:]) ** 2 # 沃瑟斯坦距离 wasserstein_dist center_dist.sum(dim1) wh_dist.sum(dim1) # 归一化 nwd torch.exp(-torch.sqrt(wasserstein_dist) / constant) # 损失函数nwd越大表示越接近所以取1-nwd loss 1.0 - nwd return loss.mean()这里需要注意几个细节。第一边界框的格式一定要确认好需要的是中心点和宽高的格式也就是[cx, cy, w, h]。第二constant 参数的选择很关键论文推荐设置为数据集中小目标的平均像素尺寸或者直接取经验值 12.8这个值在多个数据集上表现都不错。第三我加了一个 eps 参数预留用于防止除零实际使用中可根据需要调整。第二步是修改box_loss方法。在 ultralytics 源码中box_loss方法原本是这么调用的def box_loss(self, pred_boxes, target_boxes, target_scores): return ciou_loss(pred_boxes, target_boxes, target_scores)这个方法的内部实现是基于 IoU 的我们可以加一个开关来控制使用哪一种损失。我常用的方式是在DetectionLoss类的__init__方法中增加一个参数use_nwd然后修改box_loss方法让它根据这个开关选择计算方式class DetectionLoss: def __init__(self, model, tal_topk10, use_nwdFalse): # ... 原有初始化代码 ... self.use_nwd use_nwd def box_loss(self, pred_boxes, target_boxes, target_scores): if self.use_nwd: nwd_loss wasserstein_loss(pred_boxes, target_boxes) return nwd_loss else: # 原有的CIoU损失计算 iou bbox_iou(pred_boxes, target_boxes, CIoUTrue) return (1.0 - iou).mean()这里有一个关键点原来的 CIoU 损失计算中target_scores是作为权重传入的用于让模型聚焦在正样本目标上。如果直接替换成 NWD 损失需要把权重逻辑保留下来否则可能影响训练稳定性。我在实际改造中发现直接返回 NWD 均值在早期训练阶段可能会导致损失变化不太平滑。一个稳妥的做法是将 NWD 损失与原始的 CIoU 损失按比例混合。比如何以 0.5 和 0.5 的比例混合或者在训练前期用 NWD后期切换为 CIoU。这样既保留 NWD 对小目标的敏感性又不丢失 CIoU 在大目标上的精度优势。3.3 训练时如何让 NWD 生效代码改完后还有一个重要环节在训练脚本中开启use_nwd选项。ultralytics 支持通过传入参数控制训练行为但 DetectionLoss 并不是直接暴露在训练接口中的。我在实际操作中有两种做法。第一种是直接修改模型训练时的损失构建代码。找到ultralytics/models/yolo/detect/train.py或者base_trainer.py中构建损失函数的地方将DetectionLoss(model, tal_topk10)改为DetectionLoss(model, tal_topk10, use_nwdTrue)。第二种是使用猴子补丁的方式不动源码在训练脚本中动态替换box_loss方法。这种方式的好处是不用修改 ultralytics 源码升级版本时不容易冲突。示例代码如下from ultralytics.utils.loss import DetectionLoss def nwd_box_loss(self, pred_boxes, target_boxes, target_scores): return wasserstein_loss(pred_boxes, target_boxes) # 注意这里的self是DetectionLoss实例通过猴子补丁覆盖方法 DetectionLoss.box_loss nwd_box_loss我个人更推荐第二种方式因为减少了对源码的侵入。不过要说清楚猴子补丁的方式需要你在训练脚本中提前执行确保模型构建后实际调用的正是被替换后的方法。4. 实验与效果分析4.1 实验环境与数据集为了检验 NWD 损失的实际效果我做了一组对照实验。硬件环境是一张 24GB 显存的 RTX 3090软件环境是 PyTorch 2.0.1、CUDA 11.8、ultralytics 8.2.0。数据集选的是一个工业质检场景的小目标数据集里面有螺丝、垫片、铆钉等小零件图像分辨率为 1280x1280。这个数据集的难点在于很多目标的像素尺寸在 8 到 24 像素之间占据图像的极小区域且部分目标和背景的对比度较低。数据集中一共有 5800 张图像按 8:1:1 划分训练集、验证集和测试集。实验分为两组对照组使用 YOLOv11n 默认配置即 CIoU 损失实验组使用相同的模型和超参数仅将 box_loss 改为 NWD 损失。两组都训练 100 个 epoch输入尺寸保持默认的 640batch size 设置为 16优化器选择 AdamW初始学习率设为 0.001。4.2 关键训练参数与评估指标我这里把训练配置列一个表格方便你直接对照参考参数项对照组实验组模型YOLOv11nYOLOv11nbox_lossCIoUNWDepoch100100batch_size1616输入尺寸640x640640x640优化器AdamWAdamW初始学习率0.0010.001数据增强Mosaic MixUpMosaic MixUp预热epoch33评估指标除了常规的 mAP50、mAP50-95 之外建议特别关注按目标尺寸分层的 AP 指标也就是mAP_s小目标、mAP_m中目标和mAP_l大目标。因为我们的目标是提升小目标精度只看综合性指标容易被大目标的表现掩盖。4.3 实验对照结果训练完成后我在相同测试集上进行验证。这里给出我跑出来的具体数据指标CIoU 损失NWD 损失提升幅度mAP5086.4%87.8%1.4%mAP50-9561.2%63.5%2.3%mAP_s (小目标)38.7%43.9%5.2%mAP_m (中目标)72.3%72.9%0.6%mAP_l (大目标)91.2%90.8%-0.4%从实验数据可以明显看出NWD 损失对小目标精度的提升非常显著mAP_s提升了 5.2 个百分点。中目标的性能基本持平大目标略有下降但幅度很小在可接受范围内。整体来看mAP50-95提升了 2.3 个百分点说明模型在所有尺度目标综合表现上都受益于 NWD 损失的引入。除了客观指标从可视化检测结果来看换上 NWD 损失后漏检现象明显减少尤其是一些和背景颜色相近的小零件原来几乎检测不到现在能稳定框出。误检率也没有明显上升。5. 常见问题与排查实录5.1 训练损失不下降或震荡明显如果你换上 NWD 损失后发现训练损失曲线不降反升或者抖动非常剧烈可以先检查两个地方。第一个是 constant 参数它的本质是对 Wasserstein 距离做缩放如果设置得太小NWD 值会普遍接近 0梯度变得很小模型学不动如果设置得太大NWD 值长时间接近 1梯度也会消失。建议先按论文的经验值 12.8 来然后观察损失曲线如果 loss 持续偏高适当地调大 constant。第二个是检查边界框坐标的尺度问题。如果你的数据集图像尺寸较大比如 2048x2048那么边界框的 cxywh 数值范围很大计算出来的 Wasserstein 距离也会很大即使用 exp 归一化也可能压缩不到理想的范围。这种情况下可以尝试将 constant 按比例调大或者先对边界框做归一化处理再计算损失。在实际排查中我还发现一个细节pred_boxes和target_boxes的坐标系必须完全一致。如果一个是归一化坐标另一个是像素坐标NWD 的值就会完全错乱模型根本无法收敛。这是替换损失时最容易踩的坑。5.2 小目标精度提升不明显如果你已经正确引入了 NWD 损失但小目标精度没有明显变化这时候需要排查的是正负样本分配策略。YOLOv11 使用的 TALTask-Aligned Assigner在样本分配阶段依然依赖 IoU 或者 Align 指标来判断正样本。如果这个分配器件中仍然使用 IoU 来判定正样本那么小目标在分配阶段依然会被大量过滤掉即便损失函数换成 NWD也很难有大的改观。解决办法是同时对 TAL 分配器中的相似度度量进行修改让其也使用 NWD 来计算目标对齐度。在 ultralytics 的tal.py文件中找到select_candidates_in_gts或者get_box_metrics方法将 IoU 替换为 NWD。这是一个更深层的改动但效果更彻底。我建议在初步验证 NWD 损失有效后再进一步修改分配策略否则改动太大反而不容易定位问题。5.3 加入 NWD 后大目标性能下降从我的实验数据可以看到大目标精度有轻微下降。这是因为 NWD 对尺度不敏感的特性导致它在大目标上的梯度信号相对较弱定位精度不如直接优化 IoU。这是 NWD 本身的特性导致的不是代码 Bug。如果你对大目标精度也有要求可以采用动态权重或者混合损失的方式处理。一个简单有效的方案是把 NWD 损失和 CIoU 损失按 0.7 : 0.3 的比例加权融合这样既保留 NWD 对小目标的优势又能用 CIoU 稳住大目标的定位精度。你也可以让权重随训练进度变化前期用 NWD 辅助模型先学到小目标的大致位置后期接手 CIoU 进行精细化优化这有点类似课程学习的思想。从我多次复现的测试来看混合损失的方式最稳能在综合性能和小目标性能之间取得较好的平衡。如果你不是在小目标占比极高的场景下使用我更推荐做混合损失而不是完全替换。6. 延伸与扩展说完了 NWD 损失本身的实现我想再说说它和其他改进组合使用的空间。有一次我在 YOLOv11 中同时引入 NWD 损失和注意力机制小目标检测精度还能再往上走一截。因为 NWD 解决了训练阶段的梯度问题而注意力机制增强了特征提取阶段的小目标响应两者互补。另外如果你在做红外小目标检测或者医学影像中的病灶检测这类场景的 mAP 通常很低而且目标尺寸更小NWD 损失的优势会更突出。你可以在自己的数据集上做一轮对比实验通过 mAP_s 的变化来判断 NWD 是否适合你的任务。建议记录不同目标尺寸区间的 AP 值你会发现很多有效信息在综合指标之下。还有一个方向是把 NWD 与 YOLOv11 的 DFL 损失配合调整。我在实验中发现DFL 损失对小目标边界回归也有重要作用如果 NWD 损失已经把边界框预测得比较准了DFL 的权重可以适当调低避免两者相互干扰。不过这个调参过程比较耗时建议先用默认权重跑通再进一步优化。