YOLOv5改进实战 | 更换损失函数（四）之NWD（小目标检测）篇

前言

本文使用的YOLOv5版本为v7.0，该版本为YOLOv5最新版本，默认损失函数采用的是CIoU。本章节主要介绍如何将NWD损失函数应用于目标检测YOLOv5模型。

🚀🚀🚀 YOLOv5改进损失函数系列：

YOLOv5改进实战（1）| 更换损失函数（一）之EIoU、Alpha-IoU、SIoU篇
YOLOv5改进实战（2）| 更换损失函数（二）之WIOU（Wise IoU）篇
YOLOv5改进实战（3）| 更换损失函数（三）之MPDIOU（2023最新IOU）篇
YOLOv5改进实战（6）| 更换损失函数（四）之NWD（小目标检测）篇

一、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：

二、代码实现

添加损失函数

1. 在utils/metrics.py文件中添加下述源代码
   • 源代码如下：

     def wasserstein_loss(pred, target, eps=1e-7, constant=12.8):"""`Implementation of paper `Enhancing Geometric Factors intoModel Learning and Inference for Object Detection and InstanceSegmentation <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]  + epsh1 = pred[:, 3]  + epsw2 = target[:, 2] + epsh2 = target[:, 3] + epswh_distance = ((w1 - w2) ** 2 + (h1 - h2) ** 2) / 4wasserstein_2 = center_distance + wh_distancereturn torch.exp(-torch.sqrt(wasserstein_2) / constant)
     

更换NWD

按照上述更改metrics.py文件中的bbox_iou函数后，在utils/loss.py中，找到ComputeLoss类中的__call__()函数

from utils.metrics import bbox_iou, wasserstein_loss

# NWD
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)