深入解析CIoU损失函数的弊端及其优化方案

1次阅读

共计 2705 个字符，预计需要花费 7 分钟才能阅读完成。

在目标检测任务中，IoU（Intersection over Union）是最常用的评估指标之一，用于衡量预测框与真实框之间的重叠程度。为了将 IoU 直接作为损失函数进行优化，研究者们提出了 IoU 系列损失函数，包括 IoU Loss、GIoU Loss、DIoU Loss 和 CIoU Loss。其中，CIoU（Complete IoU）损失函数在 DIoU 的基础上增加了长宽比的一致性惩罚项，旨在更全面地衡量两个框的相似度。

CIoU 的设计初衷是解决 GIoU 和 DIoU 在某些情况下收敛速度慢或回归不准确的问题。具体来说，CIoU 损失函数包含了三个部分：重叠区域惩罚、中心点距离惩罚和长宽比惩罚。其数学表达式为：

$$
\mathcal{L}_{CIoU} = 1 – IoU + \frac{\rho^2(b, b^{gt})}{c^2} + \alpha v
$$

其中，$\rho$ 表示中心点之间的欧氏距离，$c$ 表示最小包围框的对角线长度，$\alpha$ 是权重系数，$v$ 用于衡量长宽比的一致性。

尽管 CIoU 在理论上比其前身更加完善，但在实际应用中仍然存在一些明显的弊端：

梯度不稳定 ：CIoU 的长宽比惩罚项在某些情况下会导致梯度不稳定，尤其是在预测框与真实框的长宽比差异较大时。这种不稳定性会影响模型的收敛速度，甚至导致训练失败。
对小目标敏感 ：CIoU 对小目标的检测性能较差。由于小目标的 IoU 值通常较小，CIoU 的长宽比惩罚项会占据主导地位，从而使得模型对小目标的回归更加困难。
长宽比惩罚项缺陷 ：CIoU 的长宽比惩罚项在某些情况下并不能有效反映两个框之间的相似度。例如，当两个框的长宽比相同但方向相反时，CIoU 的长宽比惩罚项仍然会给出较大的惩罚，这与实际情况不符。

为了克服 CIoU 的上述弊端，我们提出了一种结合 Focal Loss 和 DIoU 的改进方案。具体来说，我们首先使用 DIoU 作为基础损失函数，然后引入 Focal Loss 的思想，对难以回归的样本（如小目标）给予更大的权重。改进后的损失函数数学表达式为：

$$
\mathcal{L}_{Focal-DIoU} = \alpha_t (1 – DIoU)^\gamma
$$

其中，$\alpha_t$ 是类别权重，$\gamma$ 是调节因子，用于控制难易样本的权重分配。

以下是改进方案的 PyTorch 实现代码：

import torch
import torch.nn as nn

class FocalDIoULoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0):
        super(FocalDIoULoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(self, pred_boxes, target_boxes):
        # 计算 DIoU
        inter_area = self.compute_intersection(pred_boxes, target_boxes)
        union_area = self.compute_union(pred_boxes, target_boxes)
        iou = inter_area / union_area

        # 计算中心点距离
        center_distance = self.compute_center_distance(pred_boxes, target_boxes)
        enclosing_diagonal = self.compute_enclosing_diagonal(pred_boxes, target_boxes)
        diou = iou - (center_distance ** 2) / (enclosing_diagonal ** 2)

        # 应用 Focal Loss
        focal_diou_loss = -self.alpha * (1 - diou) ** self.gamma * torch.log(diou)
        return focal_diou_loss.mean()

    def compute_intersection(self, box1, box2):
        # 计算两个框的交集面积
        pass

    def compute_union(self, box1, box2):
        # 计算两个框的并集面积
        pass

    def compute_center_distance(self, box1, box2):
        # 计算两个框的中心点距离
        pass

    def compute_enclosing_diagonal(self, box1, box2):
        # 计算最小包围框的对角线长度
        pass

我们在 COCO 数据集上对改进前后的损失函数进行了对比实验。实验结果表明，改进后的 Focal-DIoU 损失函数在 AP（Average Precision）指标上比原始 CIoU 损失函数有显著提升，尤其是在小目标检测（AP_s）和中等目标检测（AP_m）上表现更为突出。具体数据如下表所示：