CIoU损失函数论文解析：从原理到PyTorch实战实现

1次阅读

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

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

梯度消失问题 ：当预测框与真实框没有重叠时，IoU 值为 0，无法提供有效的梯度更新方向
长宽比不敏感 ：IoU 只考虑重叠区域，对边界框的长宽比变化缺乏约束

后续提出的 GIoU（Generalized IoU）和 DIoU（Distance IoU）虽然部分解决了这些问题，但在长宽比优化上仍存在不足：

iou = intersection_area / union_area  # 传统 IoU 计算

CIoU（Complete IoU）在 DIoU 基础上增加了长宽比一致性惩罚项，其数学表达式为：

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

其中：
– $\rho$ 表示中心点欧式距离
– $c$ 是最小包围框对角线长度
– $v$ 衡量长宽比一致性
– $\alpha$ 是平衡系数

关键改进点：

中心点距离惩罚 ：通过 $\frac{\rho^2}{c^2}$ 项加速收敛
长宽比一致性项 ：
$$
v = \frac{4}{\pi^2}(\arctan\frac{w^{gt}}{h^{gt}} – \arctan\frac{w}{h})^2
$$
动态权重 ：
$$
\alpha = \frac{v}{(1-IoU)+v}
$$

import torch
import math
from torch import nn

class CIoULoss(nn.Module):
    def __init__(self, eps=1e-7):
        super(CIoULoss, self).__init__()
        self.eps = eps

    def forward(self, pred, target):
        # pred/target 格式: [x,y,w,h]
        # 转换坐标格式
        pred_xy = pred[..., :2]
        pred_wh = pred[..., 2:]

        target_xy = target[..., :2]
        target_wh = target[..., 2:]

        # 计算 IoU
        b1_x1, b1_y1 = pred_xy - pred_wh/2
        b1_x2, b1_y2 = pred_xy + pred_wh/2

        b2_x1, b2_y1 = target_xy - target_wh/2
        b2_x2, b2_y2 = target_xy + target_wh/2

        # 交集区域计算
        inter_x1 = torch.max(b1_x1, b2_x1)
        inter_y1 = torch.max(b1_y1, b2_y1)
        inter_x2 = torch.min(b1_x2, b2_x2)
        inter_y2 = torch.min(b1_y2, b2_y2)

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

        # 并集区域
        union_area = (b1_x2 - b1_x1)*(b1_y2 - b1_y1) + \
                     (b2_x2 - b2_x1)*(b2_y2 - b2_y1) - inter_area + self.eps

        iou = inter_area / union_area

        # 中心点距离惩罚项
        center_distance = torch.sum((pred_xy - target_xy)**2, dim=-1)

        # 最小包围框对角线
        enclose_x1 = torch.min(b1_x1, b2_x1)
        enclose_y1 = torch.min(b1_y1, b2_y1)
        enclose_x2 = torch.max(b1_x2, b2_x2)
        enclose_y2 = torch.max(b1_y2, b2_y2)
        c_squared = torch.sum((enclose_x2 - enclose_x1)**2 + (enclose_y2 - enclose_y1)**2, dim=-1)

        # 长宽比一致性项
        arctan = torch.atan(target_wh[...,0]/target_wh[...,1]) - \
                 torch.atan(pred_wh[...,0]/pred_wh[...,1])
        v = (4/(math.pi**2)) * torch.pow(arctan, 2)

        alpha = v / (1 - iou + v + self.eps)

        loss = 1 - iou + (center_distance / (c_squared + self.eps)) + alpha * v
        return loss.mean()