前言
本文使用的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有以下好处:
- 尺度不变性。
- 对于位置的差别变换平缓。
- 具有度量不想交的框的相似度的能力。
假设边界框 R = ( c x , c y , w , h ) R=(cx,cy,w,h) R=(cx,cy,w,h),对于两个边界框来说,其2阶Wasserstein距离可以定义为:
不过这是个距离度量,不能直接用于相似度。我们用归一化后的指数来得到一个新的度量,叫做归一化的Wasserstein距离:
基于NWD的loss:
二、代码实现
添加损失函数
- 在
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)
- 源代码如下:
更换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)