深入解析CIoU损失函数公式：从原理到PyTorch实战

1次阅读

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

目标检测任务中，边界框回归的精度直接影响模型性能。传统的 IoU（Intersection over Union）损失函数虽然直观，但在预测框和真实框没有重叠时，IoU 值为 0，梯度也为 0，导致模型无法更新参数。GIoU（Generalized IoU）虽然解决了部分问题，引入了最小闭合框的概念，但在长宽比差异较大时，仍然存在梯度消失的问题。

IoU 的局限性 ：非重叠情况下梯度为零，无法优化模型。
GIoU 的不足 ：虽然解决了非重叠问题，但对长宽比变化不敏感，梯度更新缓慢。

CIoU（Complete IoU）在 IoU 和 GIoU 的基础上，进一步引入了中心点距离和长宽比惩罚项，解决了梯度消失问题。其公式如下：

$$
CIoU = IoU – \frac{\rho^2}{c^2} – \alpha v
$$

其中：

中心点距离项 ：$\rho^2$ 是预测框和真实框中心点的欧氏距离，$c$ 是最小闭合框的对角线长度。这一项惩罚中心点偏离的情况。
长宽比惩罚项 ：$v$ 用于衡量长宽比的一致性，定义为：
$$
v = \frac{4}{\pi^2} \left(\arctan \frac{w^{gt}}{h^{gt}} – \arctan \frac{w}{h} \right)^2
$$
其中 $w^{gt}$ 和 $h^{gt}$ 是真实框的宽和高，$w$ 和 $h$ 是预测框的宽和高。
超参数 α ：用于平衡长宽比惩罚项的权重。

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

import torch
import math

def ciou_loss(pred_boxes, target_boxes):
    """
    Compute CIoU loss for bounding box regression.
    Args:
        pred_boxes: Tensor of shape (N, 4), predicted bounding boxes in (x1, y1, x2, y2) format.
        target_boxes: Tensor of shape (N, 4), ground truth bounding boxes in (x1, y1, x2, y2) format.
    Returns:
        Tensor of shape (N,), CIoU loss for each prediction.
    """
    # Convert boxes to (cx, cy, w, h) format
    pred_cx = (pred_boxes[:, 0] + pred_boxes[:, 2]) / 2
    pred_cy = (pred_boxes[:, 1] + pred_boxes[:, 3]) / 2
    pred_w = pred_boxes[:, 2] - pred_boxes[:, 0]
    pred_h = pred_boxes[:, 3] - pred_boxes[:, 1]

    target_cx = (target_boxes[:, 0] + target_boxes[:, 2]) / 2
    target_cy = (target_boxes[:, 1] + target_boxes[:, 3]) / 2
    target_w = target_boxes[:, 2] - target_boxes[:, 0]
    target_h = target_boxes[:, 3] - target_boxes[:, 1]

    # Compute IoU
    inter_x1 = torch.max(pred_boxes[:, 0], target_boxes[:, 0])
    inter_y1 = torch.max(pred_boxes[:, 1], target_boxes[:, 1])
    inter_x2 = torch.min(pred_boxes[:, 2], target_boxes[:, 2])
    inter_y2 = torch.min(pred_boxes[:, 3], target_boxes[:, 3])
    inter_area = (inter_x2 - inter_x1).clamp(min=0) * (inter_y2 - inter_y1).clamp(min=0)

    pred_area = pred_w * pred_h
    target_area = target_w * target_h
    union_area = pred_area + target_area - inter_area
    iou = inter_area / (union_area + 1e-7)

    # Compute center distance
    rho2 = (pred_cx - target_cx) ** 2 + (pred_cy - target_cy) ** 2
    c2 = (torch.max(pred_boxes[:, 2], target_boxes[:, 2]) - torch.min(pred_boxes[:, 0], target_boxes[:, 0])) ** 2 + \
         (torch.max(pred_boxes[:, 3], target_boxes[:, 3]) - torch.min(pred_boxes[:, 1], target_boxes[:, 1])) ** 2

    # Compute aspect ratio penalty
    v = (4 / (math.pi ** 2)) * (torch.atan(target_w / target_h) - torch.atan(pred_w / pred_h)) ** 2
    alpha = v / (1 - iou + v + 1e-7)

    # Compute CIoU
    ciou = iou - (rho2 / (c2 + 1e-7)) - alpha * v
    loss = 1 - ciou

    return loss

为了验证 CIoU 的效果，我们设计了不同长宽比的样本进行测试：