CIOU损失函数详解：从原理到PyTorch实战实现

1次阅读

没有评论

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

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

尺度不变性差 ：对于不同尺寸的物体，相同的 IOU 差值可能对应着完全不同的实际误差。比如大物体和小物体都差 10 像素，但小物体的 IOU 变化会更剧烈。
中心点敏感 ：当两个框没有重叠时，IOU 直接为 0，无法提供有效的梯度方向。

CIOU（Complete IOU）通过引入中心点距离惩罚和长宽比一致性惩罚，显著改善了这两个问题。

CIOU 损失可以表示为：

$$\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$ 使预测框保持与真实框相似的长宽比

import torch
import torch.nn as nn
import math

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

    def forward(self, pred, target):
        """
        Args:
            pred (tensor): [N,4] (x1,y1,x2,y2)
            target (tensor): [N,4] (x1,y1,x2,y2)
        """
        # 转换为中心点 + 宽高格式
        pred_xy = (pred[..., :2] + pred[..., 2:]) / 2
        pred_wh = (pred[..., 2:] - pred[..., :2]).clamp(min=self.eps)

        target_xy = (target[..., :2] + target[..., 2:]) / 2
        target_wh = (target[..., 2:] - target[..., :2]).clamp(min=self.eps)

        # 计算 IOU
        inter = (torch.min(pred[..., 2:], target[..., 2:]) - 
                 torch.max(pred[..., :2], target[..., :2])).clamp(0)
        union = (pred_wh.prod(-1) + target_wh.prod(-1) - inter.prod(-1))
        iou = (inter.prod(-1) + self.eps) / (union + self.eps)

        # 中心点距离
        rho2 = ((pred_xy - target_xy) ** 2).sum(-1)
        c2 = ((torch.max(pred[..., 2:], target[..., 2:]) - 
               torch.min(pred[..., :2], target[..., :2])) ** 2).sum(-1) + self.eps

        # 长宽比
        with torch.no_grad():
            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)

        return 1 - iou + (rho2 / c2) + alpha * v