YOLOv8改进实战 | IoU篇 | 更换损失函数之 NWD（Normalized Wasserstein Distance），提升小目标检测精度

---

前言

YOLOv8 官方 默认损失函数 采用的是 CIoU 。CIOU（Complete Intersection over Union）损失函数是在训练过程中用于优化边界框预测的一种度量方式。CIOU损失函数是对原始的IoU损失函数的一种改进。IoU损失函数通过计算预测框和真实框之间的交并比来进行优化，但这种简单的度量方法并没有考虑到中心点的距离以及长宽比的影响。本章节主要介绍如何将 NWD 损失函数应用于目标检测 YOLOv8 模型。

一、NWD（提升小目标检测能力）

论文链接 ： A Normalized Gaussian Wasserstein Distance for Tiny Object Detection

NWD 是一个新的度量方法来计算框和框之间的相似度，就是把框建模成高斯分布，然后用 Wasserstein 距离来度量这两个分布之间的相似度，来代替IoU。这个距离的好处是，即便是2个框完全不重叠，或者重叠很少，还是可以度量出相似度出来。另外， NWD 对于目标的尺度不敏感，对于小目标的更加的稳定。

相比于IoU，NWD有以下好处：

1. 尺度不变性。
2. 对于位置的差别变换平缓。
3. 具有度量不想交的框的相似度的能力。
   
   假设边界框 $R=（cx,cy,w,h）$ ，对于两个边界框来说，其2阶 Wasserstein 距离可以定义为：
   
   不过这是个距离度量，不能直接用于相似度。我们用归一化后的指数来得到一个新的度量，叫做归一化的Wasserstein距离：
   
   基于 NWD 的loss：

二、代码实现

项目代码

git clone -b v8.1.0 https://github.com/ultralytics/ultralytics.git

添加Loss函数

在 ultralytics/utils/metrics.py 文件中添加以下代码：

def wasserstein_loss(pred, target, eps=1e-7, constant=12.8):
    """`Implementation of paper `Enhancing Geometric Factors into
    Model Learning and Inference for Object Detection and Instance
    Segmentation <https://arxiv.org/abs/2005.03572>`_.
    Code is modified from https://github.com/Zzh-tju/CIoU.
    Args:
        pred (Tensor): Predicted bboxes of format (x_center, y_center, w, h),
            shape (n, 4).
        target (Tensor): Corresponding gt bboxes, shape (n, 4).
        eps (float): Eps to avoid log(0).
    Return:
        Tensor: Loss tensor.
    """

    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)

更换Loss函数

修改在 ultralytics/utils/loss.py 中 BboxLoss 类的 forward 函数：

"""IoU loss."""
weight = target_scores.sum(-1)[fg_mask].unsqueeze(-1)
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
# NWD
nwd = wasserstein_loss(pred_bboxes[fg_mask], target_bboxes[fg_mask])
nwd_loss = ((1.0 - nwd) * weight).sum() / target_scores_sum
iou_ratio = 0.5
loss_iou = iou_ratio * loss_iou +  (1 - iou_ratio) * nwd_loss

由于 YOLOv8 中在标签分配规则中也有用到 bbox_iou 的函数，所以需要修改 ultralytics/utils/tal.py 的 TaskAlignedAssigner 类中的 iou_calculation 函数：

def iou_calculation(self, gt_bboxes, pd_bboxes):
    """Iou calculation for horizontal bounding boxes."""
    # return bbox_iou(gt_bboxes, pd_bboxes, xywh=False, CIoU=True).squeeze(-1).clamp_(0)
    # Done: NWD
    return wasserstein_loss(gt_bboxes, pd_bboxes).squeeze(-1).clamp_(0)