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

1次阅读

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

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

当两个框没有重叠时，IoU 值为 0，梯度消失，导致模型无法学习。
IoU 损失对检测框的长宽比和中心点距离不敏感，导致回归不准确。

GIoU（Generalized IoU）虽然解决了无重叠时的梯度消失问题，但其收敛速度较慢，且对长宽比的约束仍然不足。这些局限性会直接影响模型的 mAP（mean Average Precision）指标，尤其是在 AP50 和 AP75 等高 IoU 阈值下的表现。

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

重叠面积 ：保持 IoU 的基本特性，衡量两个框的重叠程度。
中心点距离 ：惩罚预测框和真实框中心点之间的距离，公式为：
$$\mathcal{L}_{dist} = \frac{\rho^2(b, b^{gt})}{c^2}$$
其中，$\rho$ 是欧氏距离，$c$ 是最小包围框的对角线长度。
长宽比 ：约束预测框和真实框的长宽比一致性，公式为：
$$\mathcal{L}_{aspect} = \frac{v}{(1 – IoU) + v}$$
其中，$v = \frac{4}{\pi^2}(\arctan\frac{w^{gt}}{h^{gt}} – \arctan\frac{w}{h})^2$。

CIoU 的最终损失函数为：
$$\mathcal{L}{CIoU} = 1 – IoU + \mathcal{L}$$} + \mathcal{L}_{aspect

与 DIoU 和 GIoU 相比，CIoU 的收敛速度更快，尤其是在高 IoU 阈值下的表现更优。

以下是 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 损失值
    """
    # 转换格式
    box1 = box1.view(-1, 4)
    box2 = box2.view(-1, 4)

    # 计算交集面积
    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)

    # 计算并集面积
    area1 = (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1])
    area2 = (box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1])
    union_area = area1 + area2 - inter_area

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

    # 计算中心点距离
    center1_x = (box1[:, 0] + box1[:, 2]) / 2
    center1_y = (box1[:, 1] + box1[:, 3]) / 2
    center2_x = (box2[:, 0] + box2[:, 2]) / 2
    center2_y = (box2[:, 1] + box2[:, 3]) / 2
    center_dist = torch.pow(center1_x - center2_x, 2) + torch.pow(center1_y - center2_y, 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_dist = torch.pow(enclose_x2 - enclose_x1, 2) + torch.pow(enclose_y2 - enclose_y1, 2)

    # 计算长宽比
    w1 = box1[:, 2] - box1[:, 0]
    h1 = box1[:, 3] - box1[:, 1]
    w2 = box2[:, 2] - box2[:, 0]
    h2 = box2[:, 3] - box2[:, 1]
    v = (4 / (math.pi ** 2)) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
    alpha = v / (1 - iou + v + 1e-16)

    # 计算 CIoU
    ciou = iou - (center_dist / (enclose_dist + 1e-16) + alpha * v)
    return 1 - ciou

在 COCO 数据集上的实验表明，CIoU 在 AP50 和 AP75 指标上均优于 IoU 和 GIoU：