Shaun ZhouandClaude Opus 5
ef64e5bf77
- trainer: AMP、梯度裁剪、cosine/step 调度、last/best 检查点，按 monitor 指标挑选 best - losses: penalty-reduced focal loss（heatmap）+ 掩码 L1（wh/offset）+ 分类损失 （多类 BCE / 单类 CE，随 label_mode 切换） - metrics: 定位 IoU 匹配下的 P/R/F1 与 mean IoU；分类多标签 micro-F1 与单标签 accuracy， 两种模式都输出 cls_f1，故 train.monitor 无需跟着改 - optim: BatchNorm 与 bias 不做 weight decay - visualizer: SwanLab 记录损失/学习率/验证指标/显存峰值，未安装或离线时降级为警告 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
80 lines
2.8 KiB
Python
80 lines
2.8 KiB
Python
"""损失函数。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
from torch import nn
|
|
from torch.nn import functional as F
|
|
|
|
from utils.labels import MULTILABEL, check_label_mode
|
|
|
|
|
|
def gaussian_focal_loss(
|
|
pred_logits: torch.Tensor,
|
|
gt_heatmap: torch.Tensor,
|
|
alpha: float = 2.0,
|
|
beta: float = 4.0,
|
|
eps: float = 1e-4,
|
|
) -> torch.Tensor:
|
|
"""CenterNet 的 penalty-reduced focal loss，按正样本数归一化。"""
|
|
pred = pred_logits.sigmoid().clamp(eps, 1 - eps)
|
|
pos_mask = gt_heatmap.eq(1.0).float()
|
|
neg_mask = 1.0 - pos_mask
|
|
|
|
pos_loss = -((1 - pred) ** alpha) * torch.log(pred) * pos_mask
|
|
neg_loss = -((1 - gt_heatmap) ** beta) * (pred**alpha) * torch.log(1 - pred) * neg_mask
|
|
|
|
num_pos = pos_mask.sum().clamp(min=1.0)
|
|
return (pos_loss.sum() + neg_loss.sum()) / num_pos
|
|
|
|
|
|
def masked_l1_loss(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
|
|
"""只在中心点位置计算 L1。mask: (B, H, W)。"""
|
|
mask = mask.unsqueeze(1).expand_as(pred)
|
|
loss = F.l1_loss(pred * mask, target * mask, reduction="sum")
|
|
return loss / mask.sum().clamp(min=1.0)
|
|
|
|
|
|
class DetectionLoss(nn.Module):
|
|
"""heatmap + wh + offset (+ 图像级分类) 的加权和。
|
|
|
|
分类分支按 `label_mode` 切换：
|
|
multilabel（多类）—— 多热标签，逐类 BCE；
|
|
multiclass（单类）—— 互斥下标，softmax 交叉熵。
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
heatmap_weight: float = 1.0,
|
|
wh_weight: float = 0.1,
|
|
offset_weight: float = 1.0,
|
|
cls_weight: float = 1.0,
|
|
label_mode: str = MULTILABEL,
|
|
) -> None:
|
|
super().__init__()
|
|
self.label_mode = check_label_mode(label_mode)
|
|
self.heatmap_weight = heatmap_weight
|
|
self.wh_weight = wh_weight
|
|
self.offset_weight = offset_weight
|
|
self.cls_weight = cls_weight
|
|
|
|
def forward(self, outputs: dict, batch: dict) -> dict[str, torch.Tensor]:
|
|
target = batch["target"]
|
|
losses = {
|
|
"heatmap": self.heatmap_weight * gaussian_focal_loss(outputs["heatmap"], target["heatmap"]),
|
|
"wh": self.wh_weight * masked_l1_loss(outputs["wh"], target["wh"], target["mask"]),
|
|
"offset": self.offset_weight
|
|
* masked_l1_loss(outputs["offset"], target["offset"], target["mask"]),
|
|
}
|
|
if "logits" in outputs and self.cls_weight > 0:
|
|
losses["cls"] = self.cls_weight * self._classification_loss(
|
|
outputs["logits"], batch["image_labels"]
|
|
)
|
|
losses["total"] = sum(losses.values())
|
|
return losses
|
|
|
|
def _classification_loss(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
|
if self.label_mode == MULTILABEL:
|
|
return F.binary_cross_entropy_with_logits(logits, targets.float())
|
|
return F.cross_entropy(logits, targets.long())
|