优化目标检测模型：CIoU损失函数的原理与实战改进

1次阅读

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

目标检测任务中，IoU（Intersection over Union）是最基础的评估指标。但直接使用 IoU 作为损失函数时，会面临两个主要问题：

当预测框与真实框没有重叠时，IoU=0，此时梯度也为 0，无法提供有效的优化方向
IoU 无法区分不同对齐方式下的框位置差异（如中心点偏移或宽高比不符）

后续的 GIoU（Generalized IoU）通过引入最小闭合区域部分解决了第一个问题，但对框的对齐方式仍然不敏感。DIoU（Distance IoU）添加了中心点距离惩罚项，但在处理宽高比差异时效果有限。这就是 CIoU（Complete IoU）提出的背景——它同时考虑了重叠区域、中心点距离和宽高比三个因素。

CIoU 的完整定义如下：

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

其中：

$\rho$ 表示欧式距离，$b$ 和 $b^{gt}$ 分别是预测框和真实框的中心点
$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} $$

这个设计使得损失函数能够：

通过 $IoU$ 项保证基础的重叠度
通过中心点距离项加速收敛
通过宽高比项优化边界框的形状

import torch
import math

class CIoULoss:
    """
    Complete IoU loss for object detection
    Args:
        eps (float): small value to avoid division by zero
        reduction (str): 'mean' or 'sum' for loss reduction
    """def __init__(self, eps=1e-7, reduction='mean'):
        self.eps = eps
        self.reduction = reduction

    def __call__(self, pred, target):
        """
        Args:
            pred (Tensor): [N,4] (x1,y1,x2,y2) format
            target (Tensor): [N,4] (x1,y1,x2,y2) format
        Returns:
            Tensor: computed loss value
        """
        # Convert to (cx, cy, w, h) format
        pred_ctr = torch.stack([(pred[:, 0] + pred[:, 2]) / 2,
                               (pred[:, 1] + pred[:, 3]) / 2], dim=1)
        pred_wh = torch.stack([pred[:, 2] - pred[:, 0],
                              pred[:, 3] - pred[:, 1]], dim=1).clamp(min=self.eps)

        target_ctr = torch.stack([(target[:, 0] + target[:, 2]) / 2,
                                 (target[:, 1] + target[:, 3]) / 2], dim=1)
        target_wh = torch.stack([target[:, 2] - target[:, 0],
                                target[:, 3] - target[:, 1]], dim=1)

        # Calculate IoU
        inter_wh = (torch.min(pred[:, 2:], target[:, 2:]) - 
                   torch.max(pred[:, :2], target[:, :2])).clamp(min=0)
        inter_area = inter_wh[:, 0] * inter_wh[:, 1]

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

        # Center distance
        center_dist = torch.sum((pred_ctr - target_ctr)**2, dim=1)
        c = torch.sum((torch.max(pred[:, :2], target[:, :2]) - 
                      torch.min(pred[:, 2:], target[:, 2:]))**2, dim=1) + self.eps

        # Aspect ratio
        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)

        # Final loss
        loss = 1 - iou + center_dist / c + alpha * v

        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

代码特点说明：