深入解析CenterNet损失函数：从原理到实践优化

1次阅读

没有评论

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

CenterNet 作为 anchor-free 目标检测算法的代表，其核心思想是将目标检测转化为关键点（中心点）预测问题。相比传统方法，它具有三大优势：

模型结构简单，去除了复杂的 anchor 设计
推理速度快，适合实时场景
易于扩展到其他任务（如姿态估计）

但实际应用中，开发者常遇到两个典型问题：

损失函数包含多个子项，权重设置对结果影响显著但缺乏明确指导
模型收敛不稳定，容易出现预测框抖动或中心点偏移

CenterNet 的损失函数由三部分组成，用 $L$ 表示总损失：

$$ L = L_k + \lambda_{off}L_{off} + \lambda_{size}L_{size} $$

$$ L_k = \frac{-1}{N} \sum_{xyc} \begin{cases}
(1-\hat{Y}{xyc})^\alpha \log(\hat{Y}=1 \
(1-Y_{xyc})^\beta (\hat{Y}}) & \text{if} Y_{xyc{xyc})^\gamma \log(1-\hat{Y}
\end{cases} $$}) & \text{otherwise

$Y$：真实热图（高斯分布）
$\hat{Y}$：预测热图
$\alpha, \beta, \gamma$：调节难易样本权重的超参

$$ L_{off} = \frac{1}{N} \sum_{p} \left| \hat{O}_p – \left(\frac{p}{R} – \left\lfloor \frac{p}{R} \right\rfloor \right) \right| $$

$p$：真实中心点坐标
$R$：输出步长（通常为 4）
$\hat{O}_p$：预测偏移量

$$ L_{size} = \frac{1}{N} \sum_{k=1}^N \left| \hat{S}_{pk} – s_k \right| $$

$s_k$：真实目标尺寸（宽高）
$\hat{S}_{pk}$：预测尺寸

import torch
import torch.nn as nn
import torch.nn.functional as F

class CenterNetLoss(nn.Module):
    def __init__(self, alpha=2, beta=4, gamma=4, offset_weight=1, size_weight=0.1):
        super().__init__()
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        self.offset_weight = offset_weight
        self.size_weight = size_weight

    def forward(self, pred_heatmap, pred_offset, pred_size, target_heatmap, target_offset, target_size, target_mask):
        # 热图损失 (Focal Loss)
        pos_mask = target_heatmap.eq(1).float()
        neg_mask = target_heatmap.lt(1).float()

        pos_loss = torch.pow(1 - pred_heatmap, self.alpha) * torch.log(pred_heatmap + 1e-12) * pos_mask
        neg_loss = torch.pow(1 - target_heatmap, self.beta) * torch.pow(pred_heatmap, self.gamma) * \
                   torch.log(1 - pred_heatmap + 1e-12) * neg_mask

        heatmap_loss = -(pos_loss + neg_loss).sum() / max(1, pos_mask.sum())

        # 偏移量损失 (L1 Loss)
        offset_loss = F.l1_loss(pred_offset * target_mask.unsqueeze(2),
            target_offset * target_mask.unsqueeze(2),
            reduction='sum'
        ) / max(1, target_mask.sum())

        # 尺寸损失 (L1 Loss)
        size_loss = F.l1_loss(pred_size * target_mask.unsqueeze(2),
            target_size * target_mask.unsqueeze(2),
            reduction='sum'
        ) / max(1, target_mask.sum())

        total_loss = heatmap_loss + self.offset_weight * offset_loss + self.size_weight * size_loss

        return total_loss, heatmap_loss, offset_loss, size_loss