CIoU损失函数示意图解析：从理论到PyTorch实现

1次阅读

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

在目标检测任务中，边界框回归的精度直接影响模型的性能。传统的 IoU（Intersection over Union）损失函数虽然简单直观，但存在一些明显的局限性：

当两个边界框没有重叠时，IoU 值为 0，无法提供有效的梯度信息，导致训练难以收敛。
IoU 对边界框的长宽比和中心点距离不敏感，即使 IoU 相同，边界框的形状和位置可能差异很大。

GIoU（Generalized IoU）在一定程度上解决了 IoU 的问题，但仍然存在收敛速度慢和对长宽比不敏感的问题。这些问题直接影响了模型的 mAP（mean Average Precision）指标，尤其是在 AP75（IoU 阈值为 0.75）时表现更为明显。

CIoU（Complete IoU）损失函数通过引入三个核心约束项，有效解决了上述问题：

重叠面积 ：与传统 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$$

CIoU 的完整公式为：
$$\mathcal{L}{CIoU} = 1 – IoU + \mathcal{L}$$} + \alpha \mathcal{L}_{aspect

通过这三个约束项，CIoU 能够更全面地衡量边界框的相似性，从而提升模型的定位精度。

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

import torch
import math

def bbox_ciou(box1, box2):
    """
    计算两个边界框的 CIoU 损失
    :param box1: [x1, y1, x2, y2]
    :param box2: [x1, y1, x2, y2]
    :return: CIoU 损失
    """
    # 计算 IoU
    inter_x1 = torch.max(box1[0], box2[0])
    inter_y1 = torch.max(box1[1], box2[1])
    inter_x2 = torch.min(box1[2], box2[2])
    inter_y2 = torch.min(box1[3], box2[3])
    inter_area = torch.clamp(inter_x2 - inter_x1, min=0) * torch.clamp(inter_y2 - inter_y1, min=0)

    box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
    box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])
    union_area = box1_area + box2_area - inter_area
    iou = inter_area / (union_area + 1e-6)

    # 计算中心点距离
    box1_center = torch.tensor([(box1[0] + box1[2]) / 2, (box1[1] + box1[3]) / 2])
    box2_center = torch.tensor([(box2[0] + box2[2]) / 2, (box2[1] + box2[3]) / 2])
    center_distance = torch.sum((box1_center - box2_center) ** 2)

    # 计算最小包围框的对角线长度
    enclose_x1 = torch.min(box1[0], box2[0])
    enclose_y1 = torch.min(box1[1], box2[1])
    enclose_x2 = torch.max(box1[2], box2[2])
    enclose_y2 = torch.max(box1[3], box2[3])
    enclose_diagonal = (enclose_x2 - enclose_x1) ** 2 + (enclose_y2 - enclose_y1) ** 2

    # 计算长宽比
    box1_wh = torch.tensor([box1[2] - box1[0], box1[3] - box1[1]])
    box2_wh = torch.tensor([box2[2] - box2[0], box2[3] - box2[1]])
    aspect_ratio = (4 / (math.pi ** 2)) * torch.pow(torch.atan(box2_wh[0] / box2_wh[1]) - torch.atan(box1_wh[0] / box1_wh[1]), 2)

    # 计算 CIoU
    ciou = 1 - iou + (center_distance / (enclose_diagonal + 1e-6)) + aspect_ratio
    return ciou

在 COCO 数据集上的实验表明，CIoU 损失函数在 AP50 和 AP75 指标上均优于传统的 IoU 和 GIoU 损失函数。具体表现如下：