深度解析CIOU损失函数：解决目标检测中边界框回归的痛点

1次阅读

没有评论

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

在目标检测任务中，边界框回归的精度直接影响模型的性能。传统的 IOU（Intersection over Union）损失函数虽然直观，但存在两个主要问题：

梯度消失 ：当预测框与真实框没有重叠时，IOU 值为 0，导致梯度无法回传，模型无法更新参数。
对齐不准确 ：IOU 仅考虑重叠面积，忽略了中心点距离和宽高比的影响，导致对不同尺度目标的敏感度差异较大。

这些问题使得传统 IOU 损失函数在复杂场景下的表现不尽如人意。

为了解决上述问题，研究者提出了多种改进的 IOU 损失函数，包括 GIOU（Generalized IOU）、DIOU（Distance IOU）和 CIOU（Complete IOU）。以下是它们的对比：

GIOU：在 IOU 的基础上引入最小闭合框，解决了梯度消失问题，但对中心点距离和宽高比的考量不足。
$$\mathcal{L}_{GIOU} = 1 – IOU + \frac{|C – A \cup B|}{|C|}$$
DIOU：在 GIOU 的基础上增加了中心点距离的惩罚项，但对宽高比的考量仍然不足。
$$\mathcal{L}_{DIOU} = 1 – IOU + \frac{\rho^2(b, b^{gt})}{c^2}$$
CIOU：综合了中心点距离、宽高比和重叠面积，是目前最全面的改进方案。
$$\mathcal{L}_{CIOU} = 1 – IOU + \frac{\rho^2(b, b^{gt})}{c^2} + \alpha v$$
其中，$\alpha$ 是权衡参数，$v$ 是宽高比的一致性度量。

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

import torch
import math

def ciou_loss(pred_boxes, target_boxes):
    """
    Compute CIOU loss for bounding box regression.
    Args:
        pred_boxes (Tensor): Predicted bounding boxes, shape [N, 4] (x1, y1, x2, y2).
        target_boxes (Tensor): Target bounding boxes, shape [N, 4] (x1, y1, x2, y2).
    Returns:
        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]

    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 = torch.clamp(inter_x2 - inter_x1, min=0) * torch.clamp(inter_y2 - inter_y1, min=0)
    union_area = pred_w * pred_h + target_w * target_h - inter_area
    iou = inter_area / (union_area + 1e-7)

    # Compute center distance
    center_distance = (pred_cx - target_cx) ** 2 + (pred_cy - target_cy) ** 2
    c = (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
    v = (4 / (math.pi ** 2)) * torch.pow(torch.atan(target_w / (target_h + 1e-7)) - torch.atan(pred_w / (pred_h + 1e-7)), 2)
    alpha = v / (1 - iou + v + 1e-7)

    # Compute CIOU loss
    ciou = 1 - iou + center_distance / (c + 1e-7) + alpha * v
    return ciou.mean()

在 COCO 数据集上的实验表明，CIOU 损失函数在 AP（Average Precision）指标和训练收敛速度上均优于其他改进方案：