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)