学习资源站

RT-DETR改进策略【损失函数篇】 AAAI 2025 SD Loss,基于目标尺度动态调整尺度损失和位置损失的影响系数,减少 IoU 标签波动-

RT-DETR改进策略【损失函数篇】| AAAI 2025 SD Loss,基于目标尺度动态调整尺度损失和位置损失的影响系数,减少 IoU 标签波动

背景

在红外小目标检测任务里,基于交并比的损失存在波动问题,像边界框(BBox)标签的IoU损失波动最高可达86%,较小目标的IoU损失稳定性更差,严重影响模型稳定性和回归效果。

同时,现有损失函数没有充分考虑不同目标尺度下对尺度和位置的敏感度差异, 小目标的这种不稳定性更明显 ,并且掩码标签的位置损失在处理漏检物体时难以收敛,容易产生误报。

这些问题限制了模型在不同尺度目标上的检测性能,为解决这些问题,提出了 SD Loss(Scale-based Dynamic Loss) 损失函数。



MPDIoU(Intersection over Union with Minimum Points Distance) 是一种用于高效且准确的边界框回归的损失函数。

2.1 原理

SD Loss根据目标尺度动态调整Sloss和Lloss的影响系数。

在这里插入图片描述

对于BBox标签,通过计算 β B = m i n ( B g t B g t m a x × R O C × δ , δ ) \beta_{B}=min \left(\frac{B_{g t}}{B_{g t max }} × R_{O C} × \delta, \delta\right) β B = min ( B g t ma x B g t × R OC × δ , δ ) 来确定影响系数,其中 B g t B_{gt} B g t 是当前目标框面积, B g t m a x = 81 B_{gtmax}=81 B g t ma x = 81 R O C R_{OC} R OC 是原始图像与当前特征图的尺寸比例, δ \delta δ 是可调节参数。进而得到 β L B S = 1 − δ + β B \beta_{\mathcal{L}_{B S}}=1-\delta+\beta_{B} β L BS = 1 δ + β B β L B L = 1 + δ − β B \beta_{\mathcal{L}_{B L}}=1+\delta-\beta_{B} β L B L = 1 + δ β B ,最终的BBox的基于尺度的动态损失(SDB损失)为: L S D B = β L B S × L B S + β L B L × L B L \mathcal{L}_{S D B}=\beta_{\mathcal{L}_{B S}} × \mathcal{L}_{B S}+\beta_{\mathcal{L}_{B L}} × \mathcal{L}_{B L} L S D B = β L BS × L BS + β L B L × L B L

对于掩码标签,同样计算影响系数 β M = m i n ( M g t M g t m a x × R O C × δ , δ ) \beta_{M}=min \left(\frac{M_{g t}}{M_{g t max }} × R_{O C} × \delta, \delta\right) β M = min ( M g t ma x M g t × R OC × δ , δ ) ,得到 β L M S = 1 + β M \beta_{\mathcal{L}_{M S}}=1+\beta_{M} β L MS = 1 + β M β L M L = 1 − β M \beta_{\mathcal{L}_{M L}}=1-\beta_{M} β L M L = 1 β M ,掩码的基于尺度的动态损失(SDM损失)为: L S D M = β L M S × L M S + β L M L × L M L \mathcal{L}_{S D M}=\beta_{\mathcal{L}_{M S}} × \mathcal{L}_{M S}+\beta_{\mathcal{L}_{M L}} × \mathcal{L}_{M L} L S D M = β L MS × L MS + β L M L × L M L

这样,较小目标在BBox标签中对Sloss的关注权重降低,在掩码标签中增强了Sloss的影响,以此减少标签不准确对损失函数稳定性的影响。

2.2 优势

  1. 提升模型稳定性和检测精度,SD Loss通过动态调整损失系数,有效减少了IoU波动对模型的影响,增强了神经网络在不同尺度目标上的回归能力,实验结果表明,在多种模型和数据集上,使用SD Loss的模型检测精度得到显著提升。
  2. 适应性强,该损失函数在不同规模和特点的数据集(如IRSTD - 1K和SIRST - UAVB)上都能表现出良好的性能,在实际应用中,面对不同尺寸和分布的目标时具有更好的适应性。
  3. 计算简单高效,与一些涉及指数运算的损失函数(如NWD和SAFit)相比,SD Loss计算过程更为简单,不会增加过多的计算复杂度和时间成本。

论文: https://arxiv.org/pdf/2412.16986v1
源码: https://github.com/JN-Yang/PConv-SDloss-Data

三、添加步骤

3.1 ultralytics/utils/metrics.py

此处需要查看的文件是 ultralytics/utils/metrics.py

metrics.py 中定义了模型的损失函数和计算方法,我们想要加入新的损失函数就只需要将代码放到这个文件内即可

将下方的 bbox_iou 函数 替换 原本的 bbox_iou 函数即可:

def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, SDIoU=True, eps=1e-7, delta=0.5):
    """
    Calculate Intersection over Union (IoU) of box1(1, 4) to box2(n, 4).
    Args:
        box1 (torch.Tensor): A tensor representing a single bounding box with shape (1, 4).
        box2 (torch.Tensor): A tensor representing n bounding boxes with shape (n, 4).
        xywh (bool, optional): If True, input boxes are in (x, y, w, h) format. If False, input boxes are in
                               (x1, y1, x2, y2) format. Defaults to True.
        GIoU (bool, optional): If True, calculate Generalized IoU. Defaults to False.
        DIoU (bool, optional): If True, calculate Distance IoU. Defaults to False.
        CIoU (bool, optional): If True, calculate Complete IoU. Defaults to False.
        SDIoU (bool, optional): If True, calculate Scale-based Dynamic IoU. Defaults to False.
        eps (float, optional): A small value to avoid division by zero. Defaults to 1e-7.
    Returns:
        (torch.Tensor): IoU, GIoU, DIoU, or CIoU values depending on the specified flags.
    """
 
    # Get the coordinates of bounding boxes
    if xywh:  # transform from xywh to xyxy
        (x1, y1, w1, h1), (x2, y2, w2, h2) = box1.chunk(4, -1), box2.chunk(4, -1)
        w1_, h1_, w2_, h2_ = w1 / 2, h1 / 2, w2 / 2, h2 / 2
        b1_x1, b1_x2, b1_y1, b1_y2 = x1 - w1_, x1 + w1_, y1 - h1_, y1 + h1_
        b2_x1, b2_x2, b2_y1, b2_y2 = x2 - w2_, x2 + w2_, y2 - h2_, y2 + h2_
    else:  # x1, y1, x2, y2 = box1
        b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, -1)
        b2_x1, b2_y1, b2_x2, b2_y2 = box2.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
 
    # Intersection area    ∩
    inter = (b1_x2.minimum(b2_x2) - b1_x1.maximum(b2_x1)).clamp_(0) * (
        b1_y2.minimum(b2_y2) - b1_y1.maximum(b2_y1)
    ).clamp_(0)
 
    # Union Area      U
    union = w1 * h1 + w2 * h2 - inter + eps
 
    # IoU
    iou = inter / union
 
    # R_oc = 1     # The YOLO bounding box is normalized, so R_oc is equal to 1.
    if CIoU or DIoU or GIoU or SDIoU:
        cw = b1_x2.maximum(b2_x2) - b1_x1.minimum(b2_x1)  # convex (smallest enclosing box) width
        ch = b1_y2.maximum(b2_y2) - b1_y1.minimum(b2_y1)  # convex height
        if CIoU or DIoU or SDIoU:  # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
            c2 = cw**2 + ch**2 + eps  # convex diagonal squared
            rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4  # center dist ** 2
            if CIoU or SDIoU:  # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
                v = (4 / math.pi**2) * (torch.atan(w2 / h2) - torch.atan(w1 / h1)).pow(2)
                with torch.no_grad():
                    alpha = v / (v - iou + (1 + eps))
                if SDIoU:
                    beta = (w2 * h2 * delta) / 81
                    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
                    beta = beta.to(device)
                    delta = torch.tensor(delta, device=device)
                    # print(beta.device)
                    # print(torch.tensor(delta).device)
                    beta = torch.where(beta > delta, delta, beta)
                    return delta-beta + (1-delta+beta)*(iou-v*alpha) - (1+delta-beta)*(rho2/c2)  # SDIoU
                return iou - (rho2 / c2 + v * alpha)  # CIoU
            return iou - rho2 / c2  # DIoU
        c_area = cw * ch + eps  # convex area
        return iou - (c_area - union) / c_area  # GIoU https://arxiv.org/pdf/1902.09630.pdf
    return iou  # IoU

在这里插入图片描述

3.2 修改ultralytics/models/utils/loss.py

ultralytics/models/utils/loss.py 用于计算各种损失。

_get_loss_bbox函数 内修改如下代码,使模型调用此 SDIoU 损失函数。


        # loss[name_giou] = 1.0 - bbox_iou(pred_bboxes, gt_bboxes, xywh=True, GIoU=True)
        loss[name_giou] = 1.0 - bbox_iou(pred_bboxes, gt_bboxes, xywh=True, SDIoU=True)

在这里插入图片描述

此外,若训练过程中,损失为nan,则将amp设置为False。

四、成功运行截图

在这里插入图片描述