YOLO改进：【损失函数篇】提高小目标检测精度——Normalized Wasserstein Distance (NWD)损失函数

理论知识

NWD理论知识

NWD论文链接
NWD代码链接

问题背景

微小物体在许多现实世界中无处不在，微小物体的尺寸通常小于16 × 16像素，外观信息极其有限，这增加了网络模型学习鉴别微小物体特征的难度，导致微小物体检测失败率高。

IoU对不同尺度物体的灵敏度差异很大，对于6 × 6像素的微小物体，微小的位置偏差会导致明显的IoU下降(从0:53到0:06)，从而导致不准确的标签分配。而对于36 × 36像素的正常物体，IoU变化较小(从0:90到0:65)，位置偏差相同。这一现象意味着IoU度量对位置偏差离散化的对象尺度不再不变，最终导致标签分配存在以下两个缺陷。首先，IoU对微小物体的敏感性使得锚点标签的位置偏差较小，导致正负样本特征相似，网络难以收敛。其次，论文发现AI-TOD数据集中分配给每个真值(gt)的平均阳性样本数小于1，因为某些gt与任何锚点之间的IoU低于最小正阈值。因此，对于训练微小目标探测器来说，微小目标特征信息不足，无法反馈给网络，导致模型只关注于中大型目标的特征，完全忽视了对微小目标的特征信息的学习。虽然ATSS等动态分配策略可以根据物体的统计特征自适应地获得分配pos/negative标签的IoU阈值，但IoU的敏感性使得很难找到一个好的阈值并为微小物体检测器提供高质量的pos/negative样本。

NWD理论

由上述分析可以得出：IoU对于微小物体来说不是一个很好的度量，因此，文中提出了一个用Wasserstein距离度量边界盒相似性的新度量来代替标准IoU。具体来说，首先将边界框建模为二维高斯分布，然后使用提出的归一化Wasserstein距离(NWD)来度量派生高斯分布的相似性。Wasserstein距离的主要优点是，即使没有重叠或重叠可以忽略不计，它也可以度量分布的相似度。此外，NWD对不同尺度的物体不敏感，因此更适合测量微小物体之间的相似性。NWD可以应用于单级和多级锚式探测器。此外，NWD不仅可以代替标签分配中的IoU，还可以代替非最大抑制(NMS)和回归损失函数中的IoU。

具体的公式推导和数学理论就不展开介绍了，具体可以参考论文内容。

YOLO中使用NWD

在loss.py的class ComputeLoss:上面添加如下代码：

def wasserstein_loss(pred, target, eps=1e-7, constant=12.8):
    center1 = pred[:, :2]
    center2 = target[:, :2]

    whs = center1[:, :2] - center2[:, :2]

    center_distance = whs[:, 0] * whs[:, 0] + whs[:, 1] * whs[:, 1] + eps #

    w1 = pred[:, 2]  + eps
    h1 = pred[:, 3]  + eps
    w2 = target[:, 2] + eps
    h2 = target[:, 3] + eps

    wh_distance = ((w1 - w2) ** 2 + (h1 - h2) ** 2) / 4

    wasserstein_2 = center_distance + wh_distance
    return torch.exp(-torch.sqrt(wasserstein_2) / constant)

在class ComputeLoss:中的def __call__函数中，找到下面代码：

lbox += (1.0 - iou).mean()  # iou loss

# Objectness
iou = iou.detach().clamp(0).type(tobj.dtype)

将这部分代码注释调用，并用如下代码替换：

nwd = wasserstein_loss(pbox, tbox[i]).squeeze()
iou_ratio = 0.5
lbox += (1 - iou_ratio) * (1.0 - nwd).mean() + iou_ratio * (1.0 - iou).mean()  # iou loss

# Objectness
iou = (iou.detach() * iou_ratio + nwd.detach() * (1 - iou_ratio)).clamp(0, 1).type(tobj.dtype)

iou_ratio 代表IoU度量占的比例，如果数据集中微小目标占比很多，就调小该值，即增大了nwd度量。

VisDrone数据集中的实验

VisDrone数据集介绍

VisDrone数据集链接
VisDrone数据集发布于2018年，并在2019年由天津大学等进行扩展。各领域对于无人机的应用需求非常广泛。为目标检测提供了10209张图片，其中6471张图像用于训练，548张用于验证，3190张用于测试，同时还提供了96个用于目标检测的视频剪辑，包括56个用于训练(共计24201帧)，7个用于验证(共计2819帧)和33个用于测试(共计12968帧)。相对而言，VisDrone数据集场景复杂，遮挡情况严重，同时标签密集，对算法的要求也随之上升。

实验结果

下图为nwd=0.5权值的结果

下图为nwd=0.7权值的结果

下图为nwd=1.0权值的结果

下图为nwd=0.3权值的结果

由上述结果，nwd在VisDrone数据集中表现不佳。
相关训练结果下载链接：百度云：bi2d