【YOLOv8改进损失函数】基于Wasserstein Distance Loss的小目标检测评估方法（含源代码）

专栏介绍

• 攻克小目标检测难题，显著提升检测性能。
• 聚焦小目标及遮挡物检测，推动性能优化与技术创新。
• 探索工业场景下小目标检测性能提升路径与部署可行性。
• 持续更新，定期分享不同数据集上的性能提升情况 。

一、本文改进：基于Wasserstein Distance Loss的小目标检测评估方法

采用Wasserstein Distance Loss，经实测，在小目标检测任务中效果显著。

同时，推出独家原创的YOLO小目标检测创新方案，适用于YOLOv5、YOLOv8等系列YOLO模型。专栏内提供详细步骤与完整源码，助您轻松实现小目标检测的性能提升。

重点是，通过阅读本专栏，您可结合自身小目标检测数据集，在网络的Backbone、head、detect、loss等不同组件处灵活调整，实现小目标检测效果的提升与创新突破。

二、Wasserstein Distance Loss介绍

论文题目：

《A Normalized Gaussian Wasserstein Distance for Tiny Object Detection》

论文地址：

https://arxiv.org/abs/2110.13389

检测微小的物体是一个非常具有挑战性的问题，因为微小的物体只包含几个像素大小。我们证明，由于缺乏外观信息，最先进的探测器在微小物体上无法产生令人满意的结果。我们的主要观察结果是，基于 IoU 交集 （IoU） 的指标（例如 IoU 本身及其扩展）对微小物体的位置偏差非常敏感，并且在基于锚点的探测器中使用时会大大降低检测性能。为了缓解这种情况，我们提出了一种新的评估指标，使用 Wasserstein 距离进行微小物体检测。具体来说，我们首先将边界框建模为 2D 高斯分布，然后提出一个名为归一化 Wasserstein 距离 （NWD） 的新度量，通过相应的高斯分布来计算它们之间的相似性。所提出的 NWD 度量可以很容易地嵌入到任何基于锚点的检测器的赋值、非极大值抑制和损失函数中，以取代常用的 IoU 度量。我们在新的微小目标检测数据集 （AI-TOD） 上评估我们的指标，其中平均目标大小远小于现有的目标检测数据集。广泛的实验表明，当配备 NWD 指标时，我们的方法产生的性能比标准微调基线高 6.7 AP 分，比最先进的竞争对手高 6.0 AP 分。

该文档提出了一种用于小目标检测的归一化高斯Wasserstein距离（NWD），通过将边界框建模为二维高斯分布，利用Wasserstein距离计算相似性，解决了IoU对小目标位置偏差敏感的问题，显著提升了小目标检测性能。

背景与挑战

• 小目标检测难点：小目标（如AI-TOD数据集中<16×16像素）因像素少、特征不明显，现有方法（如FPN、GAN）难以有效检测。
• IoU的局限性：IoU及其扩展指标对小目标位置偏差敏感，导致标签分配不准确、训练样本不足，影响检测性能。

方法创新

1. 高斯分布建模
   • 将边界框表示为二维高斯分布，均值为框中心，协方差矩阵由宽高决定。
   • 公式：$N(\mu ,\Sigma )$，其中$\mu =[c_x,c_y{]}^T$，$\Sigma =\text{diag}(w^2/4,h^2/4)$。
2. 归一化Wasserstein距离（NWD）
   • 计算两个高斯分布的二阶Wasserstein距离：
     $W_2^2({\mathcal{N}}_a,{\mathcal{N}}_b)=\Vert m_a-m_b{\Vert }_2^2+\Vert {\Sigma }_a^{1/2}-{\Sigma }_b^{1/2}{\Vert }_F^2$
   • 归一化处理：
     $\text{NWD}=\exp \left(-\frac{\sqrt{W_2^2}}{C}\right)$
     其中C为数据集相关常数。
3. NWD的优势
   • 尺度不变性：对不同尺度目标的位置偏差不敏感。
   • 非重叠鲁棒性：可衡量无重叠或包含关系的边界框相似性。
   • 平滑性：位置偏差引起的NWD变化更平缓，利于样本区分。

应用模块

1. 标签分配
   • 用NWD替代IoU，确保每个GT框匹配足够正样本，提升训练监督质量。
2. 非极大值抑制（NMS）
   • 基于NWD的NMS更有效抑制冗余框，减少小目标漏检。
3. 回归损失函数
   • NWD损失（${\mathcal{L}}_{\text{NWD}}=1-\text{NWD}$）在无重叠或包含情况下仍提供梯度。

实验结果

1. AI-TOD数据集
   • AP提升：Faster R-CNN+NWD的AP从11.1%提升至17.8%，AP_t（8-16像素）从7.2%提升至17.0%。
   • SOTA性能：NWD-based DetectoRS达到20.8% AP，显著优于现有方法。
2. VisDrone2019数据集
   • 泛化能力：Faster R-CNN+NWD的AP_vt（2-8像素）从0.1%提升至3.8%，AP_t从6.2%提升至10.2%。
3. 消融实验
   • 模块有效性：NWD在标签分配、NMS和损失函数中均带来性能增益，其中标签分配贡献最大。

结论

• 贡献：提出NWD解决小目标检测中IoU的敏感性问题，通过高斯建模和Wasserstein距离提升检测性能。
• 应用价值：NWD可无缝集成到现有锚点检测器中，适用于多阶段和单阶段方法，显著提升小目标检测精度。

三、加入实验部分

3.1 在 ultralytics/yolo/utils/loss.py 文件中加入 Wasserstein

def Wasserstein(box1, box2, xywh=True):
    box2 = box2.T
    if xywh:
        b1_cx, b1_cy = (box1[0] + box1[2]) / 2, (box1[1] + box1[3]) / 2
        b1_w, b1_h = box1[2] - box1[0], box1[3] - box1[1]
        b2_cx, b2_cy = (box2[0] + box2[0]) / 2, (box2[1] + box2[3]) / 2
        b1_w, b1_h = box2[2] - box2[0], box2[3] - box2[1]
    else:
        b1_cx, b1_cy, b1_w, b1_h = box1[0], box1[1], box1[2], box1[3]
        b2_cx, b2_cy, b2_w, b2_h = box2[0], box2[1], box2[2], box2[3]
    cx_L2Norm = torch.pow((b1_cx - b2_cx), 2)
    cy_L2Norm = torch.pow((b1_cy - b2_cy), 2)
    p1 = cx_L2Norm + cy_L2Norm
    w_FroNorm = torch.pow((b1_w - b2_w)/2, 2)
    h_FroNorm = torch.pow((b1_h - b2_h)/2, 2)
    p2 = w_FroNorm + h_FroNorm
    return p1 + p2

class BboxLoss(nn.Module):

iou = bbox_iou(pred_bboxes[fg_mask], target_bboxes[fg_mask], xywh=False, CIoU=True)
loss_iou = ((1.0 - iou) * weight).sum() / target_scores_sum

需要替换，替换代码如下所示：

        loss_iou = 0
 
        iou = bbox_iou(pred_bboxes[fg_mask], target_bboxes[fg_mask], xywh=False, CIoU=True)
        nwd = torch.exp(
            -torch.pow(Wasserstein(pred_bboxes[fg_mask].T, target_bboxes[fg_mask], xywh=False), 1 / 2) / 1.0)
        # loss_iou = (((1.0 - iou) * weight).sum() / target_scores_sum ) * 0.5  +(((1.0 - nwd) * weight).sum() / target_scores_sum ) * 0.5
        loss_iou1 = ((1.0 - iou).mean()) * 0.5 + ((1.0 - nwd).mean()) * 0.5
 
        loss_iou = loss_iou + loss_iou1

接下来运行起来即可，下面展示epochs=5：（稳定运行！）