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

1次阅读

没有评论

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

目标检测中的边界框回归质量直接影响检测精度。传统的 IoU（Intersection over Union）损失函数存在两个致命缺陷：

当预测框与真实框无重叠时，IoU= 0 导致梯度消失
对长宽比变化不敏感，如图 1(a) 所示，不同长宽比的预测框可能具有相同 IoU 值

GIoU 通过引入最小闭包区域缓解了梯度消失问题，但仍无法解决长宽比优化问题。DIoU 增加了中心点距离惩罚项，但未考虑形状一致性。这正是 CIoU（Complete-IoU）提出的动机。

CIoU 在 DIoU 基础上增加长宽比惩罚项，完整公式为：

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

其中：
– $\rho$ 表示中心点欧氏距离
– $c$ 是最小闭包区域对角线长度
– $v$ 衡量长宽比一致性：

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

$\alpha$ 是自适应权重系数：

$$
\alpha = \frac{v}{(1-IoU)+v}
$$

图 2 展示了 CIoU 的梯度传播特性：当预测框与真实框中心点重合时，梯度主要来自长宽比惩罚项；当中心点偏离时，距离惩罚项主导梯度方向。

import torch
import torch.nn as nn
import math

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]
        # 转换坐标为 (x1,y1,x2,y2)
        pred_box = torch.cat([pred[..., :2] - pred[..., 2:]/2,
                             pred[..., :2] + pred[..., 2:]/2], dim=-1)

        target_box = torch.cat([target[..., :2] - target[..., 2:]/2,
                               target[..., :2] + target[..., 2:]/2], dim=-1)

        # 计算 IoU
        inter_area = self._calculate_inter_area(pred_box, target_box)
        union_area = self._calculate_union_area(pred_box, target_box)
        iou = inter_area / (union_area + self.eps)

        # 中心点距离惩罚项
        center_distance = torch.sum(torch.pow(pred[..., :2] - target[..., :2], 2), dim=-1)
        c_diagonal = self._calculate_enclosing_diagonal(pred_box, target_box)
        diou_term = center_distance / (c_diagonal + self.eps)

        # 长宽比惩罚项
        arctan = torch.atan(target[..., 2]/target[..., 3]) - torch.atan(pred[..., 2]/pred[..., 3])
        v = (4 / (math.pi ** 2)) * torch.pow(arctan, 2)
        alpha = v / (1 - iou + v + self.eps)

        return 1 - iou + diou_term + alpha * v

    def _calculate_inter_area(self, box1, box2):
        # 实现交集计算...
        pass

    def _calculate_union_area(self, box1, box2):
        # 实现并集计算...
        pass

    def _calculate_enclosing_diagonal(self, box1, box2):
        # 计算最小闭包区域对角线...
        pass

测试环境：RTX 3090, PyTorch 1.12, YOLOv5s 模型