CIoU损失函数论文解析：从理论到PyTorch实战实现

1次阅读

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

在目标检测任务中，边界框回归的精度直接影响模型的性能。传统 IoU（Intersection over Union）损失函数存在两个主要缺陷：当预测框与真实框没有重叠时，IoU 值为 0，导致梯度消失；此外，IoU 无法区分不同对齐方式的重叠情况，导致收敛方向不稳定。GIoU（Generalized IoU）虽然解决了非重叠情况下的梯度消失问题，但在边界框包含情况下仍存在收敛速度慢的问题。

CIoU（Complete IoU）损失函数通过引入三个关键改进，有效解决了上述问题：

中心点距离惩罚项：通过最小化预测框与真实框中心点之间的欧氏距离，加速收敛。

$$\mathcal{L}_{dist} = \frac{\rho^2(b, b^{gt})}{c^2}$$

其中，$\rho$ 表示中心点距离，$c$ 是最小包围框的对角线长度。

宽高比一致性约束：通过惩罚宽高比差异，提高边界框的匹配精度。

$$\mathcal{L}_{aspect} = \frac{4}{\pi^2}(\arctan\frac{w^{gt}}{h^{gt}} – \arctan\frac{w}{h})^2$$

重叠区域权重：通过动态调整重叠区域权重，优化梯度传播。

$$\mathcal{L}{CIoU} = 1 – IoU + \mathcal{L}$$} + \alpha \mathcal{L}_{aspect

以下是完整的 CIoU 损失函数的 PyTorch 实现代码：

import torch
import math

class CIoULoss(torch.nn.Module):
    def __init__(self):
        super(CIoULoss, self).__init__()

    def forward(self, pred, target):
        """
        Args:
            pred (Tensor): [N, 4] (x1, y1, x2, y2)
            target (Tensor): [N, 4] (x1, y1, x2, y2)
        Returns:
            loss (Tensor): scalar
        """
        # 确保输入坐标在 [0,1] 范围内
        pred = torch.clamp(pred, 0, 1)
        target = torch.clamp(target, 0, 1)

        # 计算交集区域
        inter_x1 = torch.max(pred[:, 0], target[:, 0])
        inter_y1 = torch.max(pred[:, 1], target[:, 1])
        inter_x2 = torch.min(pred[:, 2], target[:, 2])
        inter_y2 = torch.min(pred[:, 3], target[:, 3])

        inter_area = torch.clamp(inter_x2 - inter_x1, min=0) * torch.clamp(inter_y2 - inter_y1, min=0)

        # 计算并集区域
        pred_area = (pred[:, 2] - pred[:, 0]) * (pred[:, 3] - pred[:, 1])
        target_area = (target[:, 2] - target[:, 0]) * (target[:, 3] - target[:, 1])
        union_area = pred_area + target_area - inter_area

        # 计算 IoU
        iou = inter_area / (union_area + 1e-7)

        # 计算中心点距离
        pred_center = torch.stack([(pred[:, 0] + pred[:, 2]) / 2, (pred[:, 1] + pred[:, 3]) / 2], dim=1)
        target_center = torch.stack([(target[:, 0] + target[:, 2]) / 2, (target[:, 1] + target[:, 3]) / 2], dim=1)
        center_distance = torch.sum((pred_center - target_center) ** 2, dim=1)

        # 计算最小包围框对角线长度
        enclose_x1 = torch.min(pred[:, 0], target[:, 0])
        enclose_y1 = torch.min(pred[:, 1], target[:, 1])
        enclose_x2 = torch.max(pred[:, 2], target[:, 2])
        enclose_y2 = torch.max(pred[:, 3], target[:, 3])
        c_squared = torch.sum((enclose_x2 - enclose_x1) ** 2 + (enclose_y2 - enclose_y1) ** 2, dim=0)

        # 计算宽高比一致性
        pred_wh = pred[:, 2:] - pred[:, :2]
        target_wh = target[:, 2:] - target[:, :2]
        v = (4 / (math.pi ** 2)) * torch.pow(torch.atan(pred_wh[:, 0] / (pred_wh[:, 1] + 1e-7)) - 
                                            torch.atan(target_wh[:, 0] / (target_wh[:, 1] + 1e-7)), 2)
        alpha = v / (1 - iou + v + 1e-7)

        # 计算 CIoU 损失
        loss = 1 - iou + (center_distance / (c_squared + 1e-7)) + alpha * v

        return loss.mean()

我们在 COCO 数据集上进行了对比实验，测试环境为 NVIDIA V100 GPU，PyTorch 1.7.1。实验结果如下：