深入解析CIoU损失函数原理图：从理论到PyTorch实现

1次阅读

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

在目标检测任务中，边界框回归（Bounding Box Regression）是决定检测精度的关键环节。传统的 IoU（Intersection over Union）指标虽然直观，但在非重叠或部分重叠情况下存在梯度消失问题。随后出现的 GIoU（Generalized IoU）和 DIoU（Distance IoU）逐步解决了部分问题，但仍未全面考虑长宽比的一致性。

CIoU（Complete IoU）在 DIoU 基础上引入长宽比一致性惩罚项，其数学定义为：

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

其中：
– $\rho$ 表示预测框与真实框中心点的欧氏距离
– $c$ 是最小包围矩形的对角线长度
– $v$ 衡量长宽比一致性：
$$
v = \frac{4}{\pi^2}(\arctan\frac{w^{gt}}{h^{gt}} – \arctan\frac{w}{h})^2
$$
– $\alpha$ 是权重系数：
$$
\alpha = \frac{v}{(1-IoU)+v}
$$

在 COCO val2017 数据集上，使用 YOLOv3 框架的对比结果：

损失函数	AP@0.5	AP@0.75
IoU	58.2	34.1
GIoU	60.1	37.8
DIoU	61.4	39.2
CIoU	62.8	40.5

import torch
import math

def bbox_ciou(box1: torch.Tensor, box2: torch.Tensor) -> torch.Tensor:
    """
    计算 CIoU 损失
    :param box1: 预测框 [x1,y1,x2,y2] 格式
    :param box2: 真实框 [x1,y1,x2,y2] 格式
    :return: CIoU 损失值
    """
    # 转换为中心点 + 宽高表示
    b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, dim=-1)
    b2_x1, b2_y1, b2_x2, b2_y2 = box2.chunk(4, dim=-1)
    w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1
    w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1

    # 计算 IoU
    inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
            (torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)
    union = w1 * h1 + w2 * h2 - inter
    iou = inter / (union + 1e-7)

    # 中心点距离
    c_x1, c_y1 = torch.min(b1_x1, b2_x1), torch.min(b1_y1, b2_y1)
    c_x2, c_y2 = torch.max(b1_x2, b2_x2), torch.max(b1_y2, b2_y2)
    c_diag = ((c_x2 - c_x1) ** 2 + (c_y2 - c_y1) ** 2) ** 0.5

    # 中心点欧氏距离
    rho = ((b1_x1 + b1_x2 - b2_x1 - b2_x2) ** 2 / 4 + 
           (b1_y1 + b1_y2 - b2_y1 - b2_y2) ** 2 / 4) ** 0.5

    # 长宽比惩罚项
    v = (4 / (math.pi ** 2)) * torch.pow(torch.atan(w2 / (h2 + 1e-7)) - torch.atan(w1 / (h1 + 1e-7)), 2)
    alpha = v / (1 - iou + v + 1e-7)

    return 1 - iou + (rho ** 2) / (c_diag ** 2 + 1e-7) + alpha * v