当前位置：首页 > news >正文

YOLOv11小白的进击之路（九）创新YOLO11损失函数之NWD损失函数源码解读

news 来源：原创 2025/5/18 6:05:09

之前的博客也有对YOLO11的损失函数进行过源码分析，可以参考：YOLOv11小白的进击之路（六）创新YOLO的iou及损失函数时的源码分析_yolov11的损失函数是什么-CSDN博客最近在做小目标检测的时候注意到了NWD损失函数，这里对其源码进行分析，结合YOLO11更改损失函数后在我自己的数据集上有一定的提升，但是效果并不是很明显（可能我的数据集目标尺寸并没有特别小？）这里留待之后讨论。

NWD

NWD是一种用于目标检测的损失函数，通过将边界框建模为高斯分布，并计算两个分布之间的Wasserstein距离来衡量相似性。具体步骤包括：

边界框参数化：将边界框 R=(cx,cy,w,h)（中心点坐标和宽高）转换为高斯分布，其中均值为中心点 (cx,cy)，协方差矩阵由宽高 (w,h)决定。
Wasserstein距离计算：衡量两个分布之间的差异，公式为 W2=∥c1−c2∥2+∥w1−w2∥2+∥h1−h2∥2W2=∥c1−c2∥2+∥w1−w2∥2+∥h1−h2∥2，其中 c为中心点，w,h 为宽高。
归一化相似度：通过指数归一化将距离转换为相似度值 NWD=exp⁡(−W/C)，其中 C常数。

优势

NWD相较于传统IoU（交并比）的改进主要体现在：

对小目标更友好：对小尺度目标的检测更稳定，避免因像素级偏差导致的IoU剧烈波动。
处理不重叠框的能力：即使两个边界框完全不重叠，NWD仍能有效度量相似性，而IoU此时为0，无法提供梯度。
尺度不变性：对目标尺寸变化不敏感，更适合多尺度检测任务。

应用场景

小目标检测：如卫星图像、医学影像中的微小物体检测。
密集场景检测：在目标重叠或遮挡严重时，NWD能更平滑地优化模型

源码分析

def wasserstein_loss(pred, target, eps=1e-7, constant=12.8):
    r"""`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_min, y_min, x_max, y_max),
            shape (n, 4).
        target (Tensor): Corresponding gt bboxes, shape (n, 4).
        eps (float): Eps to avoid log(0).
    Return:
        Tensor: Loss tensor.
    """

    b1_x1, b1_y1, b1_x2, b1_y2 = pred.chunk(4, -1)
    b2_x1, b2_y1, b2_x2, b2_y2 = target.chunk(4, -1)
    w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
    w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
    
    b1_x_center, b1_y_center = b1_x1 + w1 / 2, b1_y1 + h1 / 2
    b2_x_center, b2_y_center = b2_x1 + w2 / 2, b2_y1 + h2 / 2
    center_distance = (b1_x_center - b2_x_center) ** 2 + (b1_y_center - b2_y_center) ** 2 + eps
    wh_distance = ((w1 - w2) ** 2 + (h1 - h2) ** 2) / 4

    wasserstein_2 = center_distance + wh_distance
    return torch.exp(-torch.sqrt(wasserstein_2) / constant)

我们现将参数分解：

b1_x1, b1_y1, b1_x2, b1_y2 = pred.chunk(4, -1)
b2_x1, b2_y1, b2_x2, b2_y2 = target.chunk(4, -1)

目的：将预测框（pred）和真实框（target）的坐标拆分成左上角（x1, y1）和右下角（x2, y2）。
chunk(4, -1)：沿最后一个维度（即列方向）将每个框的4个坐标拆分成4个变量。例如，pred的格式是(x_min, y_min, x_max, y_max)，拆分成4个张量后分别对应这四个坐标值。

再计算框的宽高：

w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps

目的：计算预测框和真实框的宽度（w）和高度（h）。
+ eps：为了防止宽高为0时出现数值不稳定（例如除零错误），加上一个极小值eps（默认为1e-7）。

计算中心点坐标：

b1_x_center, b1_y_center = b1_x1 + w1 / 2, b1_y1 + h1 / 2
b2_x_center, b2_y_center = b2_x1 + w2 / 2, b2_y1 + h2 / 2

目的：通过左上角坐标和宽高，计算预测框和真实框的中心点坐标。

计算中心点距离：

center_distance = (b1_x_center - b2_x_center)**2 + (b1_y_center - b2_y_center)**2 + eps

目的：计算两个框中心点的欧氏距离的平方。

计算宽高距离：

wh_distance = ((w1 - w2)**2 + (h1 - h2)**2) / 4

目的：计算两个框宽高差异的平方，并除以4（对应高斯分布协方差矩阵的推导）。

合并得到Wasserstein距离：

wasserstein_2 = center_distance + wh_distance

目的：将中心点距离和宽高距离相加，得到最终的Wasserstein距离平方。
数学依据：
将边界框建模为高斯分布后，两个分布的Wasserstein距离平方可以分解为：

其中协方差矩阵 ΣΣ 由宽高决定（）。

最后归一化损失值：

return torch.exp(-torch.sqrt(wasserstein_2) / constant)

目的：将Wasserstein距离转换为相似度值（取值范围0~1），作为损失函数。
步骤：
1. 开平方：torch.sqrt(wasserstein_2) 得到真正的Wasserstein距离 WW。
2. 归一化：通过指数函数归一化，公式为 NWD=exp⁡(−W/C)，其中 C 是归一化常数
3. 损失值：NWD越大，表示两个框越相似，因此损失函数需要最小化 1−NWD。这里我们直接返回相似度值，实际使用时可以用 1 - NWD 作为损失。