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

1次阅读

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

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

梯度消失问题：当预测框与真实框无重叠时，IoU= 0 导致梯度为 0，无法更新参数。数学表达式为：
$$L_{IoU} = 1 – \frac{|A \cap B|}{|A \cup B|}$$
其中 $A$ 和 $B$ 分别表示预测框和真实框的面积。
尺度敏感性：IoU 对长宽比变化不敏感，例如两个相同 IoU 的预测框可能长宽比差异很大。这源于边界框参数（中心点 xy 与宽高 wh）的耦合性，用公式表示为：
$$\nabla_{xy}L = f(wh), \quad \nabla_{wh}L = f(xy)$$

CIOU（Complete-IoU）通过引入三项惩罚项来解决上述问题：

中心点距离惩罚：
$$\rho^2(b,b^{gt}) = (x-x^{gt})^2 + (y-y^{gt})^2$$
其中 $(x,y)$ 和 $(x^{gt},y^{gt})$ 分别是预测框和真实框的中心坐标。
长宽比一致性惩罚：
$$\alpha = \frac{v}{(1-IoU)+v}, \quad v = \frac{4}{\pi^2}(\arctan\frac{w^{gt}}{h^{gt}} – \arctan\frac{w}{h})^2$$
该设计使得长宽比差异大的预测框会受到更大惩罚。
完整公式：
$$L_{CIoU} = 1 – IoU + \frac{\rho^2(b,b^{gt})}{c^2} + \alpha v$$
其中 $c$ 是最小包围框的对角线长度。

CIOU 的梯度计算避免了零梯度问题。以中心点 x 坐标为例：
$$\frac{\partial L}{\partial x} = \frac{2(x-x^{gt})}{c^2} + \text{IoU 相关项}$$
即使 IoU= 0 时，第一项仍能提供有效梯度。

import torch

def bbox_ciou(pred_boxes, target_boxes, eps=1e-7):
    """
    pred_boxes: Tensor[N,4] (x1,y1,x2,y2)
    target_boxes: Tensor[N,4]
    """
    # 转换为中心点 + 宽高格式
    pred_ctr = (pred_boxes[:, :2] + pred_boxes[:, 2:]) / 2
    pred_wh = pred_boxes[:, 2:] - pred_boxes[:, :2]
    target_ctr = (target_boxes[:, :2] + target_boxes[:, 2:]) / 2
    target_wh = target_boxes[:, 2:] - target_boxes[:, :2]

    # IoU 计算（数值稳定性保护）inter_upleft = torch.max(pred_boxes[:, :2], target_boxes[:, :2])
    inter_botright = torch.min(pred_boxes[:, 2:], target_boxes[:, 2:])
    inter_wh = (inter_botright - inter_upleft).clamp(min=0)
    inter_area = inter_wh[:, 0] * inter_wh[:, 1]

    union_area = pred_wh[:,0]*pred_wh[:,1] + target_wh[:,0]*target_wh[:,1] - inter_area + eps
    iou = inter_area / union_area

    # 中心点距离惩罚
    center_dist = torch.sum((pred_ctr - target_ctr)**2, dim=1)
    c = torch.sum((torch.max(pred_boxes[:, 2:], target_boxes[:, 2:]) - 
                   torch.min(pred_boxes[:, :2], target_boxes[:, :2]))**2, dim=1)
    rho2 = center_dist / (c + eps)

    # 长宽比惩罚
    atan_diff = torch.atan(target_wh[:,0]/target_wh[:,1]) - torch.atan(pred_wh[:,0]/pred_wh[:,1])
    v = 4*(atan_diff**2) / (math.pi**2)
    alpha = v / (1 - iou + v + eps)

    return 1 - iou + rho2 + alpha*v

关键实现技巧：
– 使用 clamp(min=0) 处理无重叠情况
– 添加微小值 eps 防止除零错误
– 向量化计算所有框对的 CIOU

在 COCO val2017 上的测试结果（YOLOv5s 模型）：