CIoU损失函数原理剖析与PyTorch实战：解决目标检测中的边界框回归难题

1次阅读

没有评论

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

在目标检测任务中，边界框（Bounding Box）回归的精度直接影响模型性能。传统的 IoU（Intersection over Union）损失函数存在以下问题：

梯度消失 ：当预测框与真实框没有重叠时，IoU 为零，无法提供有效的梯度更新。
尺度不变性 ：IoU 对目标尺度敏感，不同尺度的目标可能产生相同的 IoU 值。
宽高比敏感 ：IoU 无法区分不同宽高比的边界框，导致回归精度下降。

CIoU（Complete IoU）损失函数在 DIoU 的基础上引入了宽高比的一致性度量，其公式如下：

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

其中：

$\rho$ 是欧氏距离，$b$ 和 $b^{gt}$ 分别是预测框和真实框的中心点。
$c$ 是最小外接矩形的对角线长度。
$\alpha$ 是权重系数，$v$ 是宽高比一致性度量：

$$v = \frac{4}{\pi^2} (\arctan\frac{w^{gt}}{h^{gt}} – \arctan\frac{w}{h})^2$$

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

import torch
import math

def bbox_ciou(pred_boxes: torch.Tensor, gt_boxes: torch.Tensor, eps: float = 1e-7) -> torch.Tensor:
    """
    Compute CIoU loss between predicted and ground truth boxes.

    Args:
        pred_boxes (torch.Tensor): Predicted bounding boxes (x1, y1, x2, y2).
        gt_boxes (torch.Tensor): Ground truth bounding boxes (x1, y1, x2, y2).
        eps (float): Small constant for numerical stability.

    Returns:
        torch.Tensor: CIoU loss.
    """
    # 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]

    gt_cx = (gt_boxes[..., 0] + gt_boxes[..., 2]) / 2
    gt_cy = (gt_boxes[..., 1] + gt_boxes[..., 3]) / 2
    gt_w = gt_boxes[..., 2] - gt_boxes[..., 0]
    gt_h = gt_boxes[..., 3] - gt_boxes[..., 1]

    # Calculate intersection area
    inter_x1 = torch.max(pred_boxes[..., 0], gt_boxes[..., 0])
    inter_y1 = torch.max(pred_boxes[..., 1], gt_boxes[..., 1])
    inter_x2 = torch.min(pred_boxes[..., 2], gt_boxes[..., 2])
    inter_y2 = torch.min(pred_boxes[..., 3], gt_boxes[..., 3])

    inter_area = torch.clamp(inter_x2 - inter_x1, min=0) * torch.clamp(inter_y2 - inter_y1, min=0)

    # Calculate union area
    pred_area = pred_w * pred_h
    gt_area = gt_w * gt_h
    union_area = pred_area + gt_area - inter_area + eps

    # Calculate IoU
    iou = inter_area / union_area

    # Calculate center distance
    center_distance = (pred_cx - gt_cx) ** 2 + (pred_cy - gt_cy) ** 2

    # Calculate enclosing box diagonal distance
    enclose_x1 = torch.min(pred_boxes[..., 0], gt_boxes[..., 0])
    enclose_y1 = torch.min(pred_boxes[..., 1], gt_boxes[..., 1])
    enclose_x2 = torch.max(pred_boxes[..., 2], gt_boxes[..., 2])
    enclose_y2 = torch.max(pred_boxes[..., 3], gt_boxes[..., 3])

    enclose_distance = (enclose_x2 - enclose_x1) ** 2 + (enclose_y2 - enclose_y1) ** 2

    # Calculate aspect ratio consistency
    v = (4 / (math.pi ** 2)) * torch.pow(torch.atan(gt_w / gt_h) - torch.atan(pred_w / pred_h), 2)
    alpha = v / (1 - iou + v + eps)

    # Combine all components
    ciou = iou - (center_distance / (enclose_distance + eps)) - (alpha * v)
    loss = 1 - ciou

    return loss.mean()

在 COCO 数据集上的实验表明，CIoU 损失函数相比 IoU、GIoU 和 DIoU，能够提升 mAP 约 3%-5%。以下是一些调参建议：