yolov7添加Normalized Wasserstein Distance (NWD)损失函数

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

01 在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)

02 在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和原始iou的权重。如果数据集中微小目标占比很多，就调小该值，即增大了nwd度量。调节iou_ratio找到最优的设置。

最后：

如果你想要进一步了解更多的相关知识，可以关注下面公众号联系~会不定期发布相关设计内容包括但不限于如下内容:信号处理、通信仿真、算法设计、matlab appdesigner，gui设计、simulink仿真......希望能帮到你！