深入解析CIOU损失函数原理：从数学推导到PyTorch实现

1次阅读

没有评论

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

在目标检测任务中，边界框回归是核心环节。传统的 IoU（Intersection over Union）损失函数存在两个主要缺陷：

梯度消失问题：当预测框与真实框不相交时，IoU 值为 0，此时无法计算梯度，导致模型无法更新参数。数学表达式为：

$$\mathcal{L}_{IoU} = 1 – \frac{|A \cap B|}{|A \cup B|}$$

其中 (A) 和(B)分别表示预测框和真实框的区域。

尺度敏感性：IoU 对边界框的长宽比变化不够敏感。例如，两个长宽比差异很大的框可能具有相同的 IoU 值，但实际回归效果差异很大。

CIOU（Complete-IoU）在 DIoU 的基础上增加了长宽比一致性惩罚项，其公式为：

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

其中：

(\rho^2(b, b^{gt}))表示预测框中心点与真实框中心点的欧氏距离
(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}$$

CIOU 通过引入中心点距离和长宽比惩罚项，确保在预测框与真实框不相交时仍能产生有效梯度。以中心点坐标 (x) 为例，其梯度为：

$$\frac{\partial \mathcal{L}_{CIoU}}{\partial x} = \frac{2(x – x^{gt})}{c^2} + \text{其他项}$$

这说明即使 IoU 为 0，中心点距离项仍能提供梯度信号。

以下是 PyTorch 实现的 CIOU 损失函数：

import torch
import math

def bbox_ciou(boxes1, boxes2):
    """
    计算 CIOU 损失
    Args:
        boxes1: 预测框 [..., 4] (x1,y1,x2,y2)
        boxes2: 真实框 [..., 4] (x1,y1,x2,y2)
    """
    # 转换为中心点 + 宽高表示
    boxes1_xy = (boxes1[..., :2] + boxes1[..., 2:]) / 2
    boxes1_wh = boxes1[..., 2:] - boxes1[..., :2]
    boxes2_xy = (boxes2[..., :2] + boxes2[..., 2:]) / 2
    boxes2_wh = boxes2[..., 2:] - boxes2[..., :2]

    # 计算 IoU
    inter_area = torch.prod(torch.minimum(boxes1_wh, boxes2_wh).clamp(min=0), 
        dim=-1
    )
    union_area = torch.prod(boxes1_wh, dim=-1) + torch.prod(boxes2_wh, dim=-1) - inter_area
    iou = inter_area / (union_area + 1e-7)  # 数值稳定性保护

    # 中心点距离
    center_distance = torch.sum(torch.pow(boxes1_xy - boxes2_xy, 2), dim=-1)

    # 最小包围框对角线长度
    enclose_wh = torch.maximum(boxes1_wh, boxes2_wh)
    enclose_diagonal = torch.sum(torch.pow(enclose_wh, 2), dim=-1)

    # 长宽比一致性
    arctan = torch.atan(boxes1_wh[..., 0] / (boxes1_wh[..., 1] + 1e-7)) - \
             torch.atan(boxes2_wh[..., 0] / (boxes2_wh[..., 1] + 1e-7))
    v = 4 * torch.pow(arctan / math.pi, 2)
    alpha = v / (1 - iou + v + 1e-7)

    # 组合 CIOU 损失
    ciou = 1 - iou + center_distance / (enclose_diagonal + 1e-7) + alpha * v
    return ciou.mean()

在 COCO 数据集上的对比实验结果：