feat(engine): 训练评估循环、CenterNet 损失与指标

- 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>

engine/__init__.py

@@ -0,0 +1,17 @@
+"""训练 / 评估流程。"""
+
+from engine.checkpoint import load_checkpoint, save_checkpoint
+from engine.evaluator import Evaluator
+from engine.losses import DetectionLoss
+from engine.optim import build_optimizer, build_scheduler
+from engine.trainer import Trainer
+
+__all__ = [
+    "Trainer",
+    "Evaluator",
+    "DetectionLoss",
+    "build_optimizer",
+    "build_scheduler",
+    "save_checkpoint",
+    "load_checkpoint",
+]

engine/checkpoint.py

@@ -0,0 +1,46 @@
+"""检查点存取。"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import torch
+from torch import nn
+
+
+def save_checkpoint(
+    path: str | Path,
+    model: nn.Module,
+    optimizer: torch.optim.Optimizer | None = None,
+    epoch: int = 0,
+    best_metric: float | None = None,
+    config: dict | None = None,
+) -> Path:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(
+        {
+            "model": model.state_dict(),
+            "optimizer": optimizer.state_dict() if optimizer else None,
+            "epoch": epoch,
+            "best_metric": best_metric,
+            "config": config,
+        },
+        path,
+    )
+    return path
+
+
+def load_checkpoint(
+    path: str | Path,
+    model: nn.Module,
+    optimizer: torch.optim.Optimizer | None = None,
+    map_location: Any = "cpu",
+    strict: bool = True,
+) -> dict:
+    state = torch.load(Path(path), map_location=map_location, weights_only=False)
+    model.load_state_dict(state["model"], strict=strict)
+    if optimizer is not None and state.get("optimizer"):
+        optimizer.load_state_dict(state["optimizer"])
+    return state

engine/evaluator.py

@@ -0,0 +1,87 @@
+"""评估：跑一遍验证集，输出分类与定位指标。"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+
+from engine.metrics import AverageMeter, DetectionMetric, MultiLabelMetric, SingleLabelMetric
+from models.postprocess import decode_heatmap
+from utils.labels import MULTILABEL
+from utils.logging import get_logger
+
+logger = get_logger("engine.evaluator")
+
+
+class Evaluator:
+    def __init__(
+        self,
+        model: nn.Module,
+        dataloader: DataLoader,
+        device: torch.device,
+        criterion: nn.Module | None = None,
+        score_threshold: float = 0.3,
+        iou_threshold: float = 0.5,
+        topk: int = 50,
+    ) -> None:
+        self.model = model
+        self.dataloader = dataloader
+        self.device = device
+        self.criterion = criterion
+        self.score_threshold = score_threshold
+        self.topk = topk
+        self.det_metric = DetectionMetric(iou_threshold)
+        # 分类指标按模型的 label_mode 选择：多类用多标签指标，单类用互斥分类指标。
+        self.label_mode = getattr(model, "label_mode", MULTILABEL)
+        if self.label_mode == MULTILABEL:
+            self.cls_metric: MultiLabelMetric | SingleLabelMetric = MultiLabelMetric()
+        else:
+            num_outputs = model.cls_head.num_outputs if model.cls_head is not None else 1
+            self.cls_metric = SingleLabelMetric(num_outputs)
+
+    @torch.no_grad()
+    def evaluate(self) -> dict[str, float]:
+        self.model.eval()
+        self.det_metric.reset()
+        self.cls_metric.reset()
+        loss_meter = AverageMeter()
+
+        stride = getattr(self.model, "out_stride", 4)
+        for batch in self.dataloader:
+            images = batch["images"].to(self.device, non_blocking=True)
+            outputs = self.model(images)
+
+            if self.criterion is not None:
+                batch_on_device = _to_device(batch, self.device)
+                losses = self.criterion(outputs, batch_on_device)
+                loss_meter.update(float(losses["total"]), images.shape[0])
+
+            detections = decode_heatmap(
+                outputs["heatmap"],
+                outputs["wh"],
+                outputs["offset"],
+                stride=stride,
+                topk=self.topk,
+                score_threshold=self.score_threshold,
+            )
+            for i, det in enumerate(detections):
+                self.det_metric.update(det["boxes"], det["labels"], batch["boxes"][i], batch["labels"][i])
+
+            if "logits" in outputs:
+                self.cls_metric.update(outputs["logits"], batch["image_labels"])
+
+        metrics: dict[str, float] = {**self.det_metric.compute(), **self.cls_metric.compute()}
+        if self.criterion is not None:
+            metrics["loss"] = loss_meter.avg
+        return metrics
+
+
+def _to_device(batch: dict, device: torch.device) -> dict:
+    moved: dict[str, Any] = dict(batch)
+    moved["images"] = batch["images"].to(device, non_blocking=True)
+    moved["target"] = {k: v.to(device, non_blocking=True) for k, v in batch["target"].items()}
+    moved["image_labels"] = batch["image_labels"].to(device, non_blocking=True)
+    return moved

engine/losses.py

@@ -0,0 +1,79 @@
+"""损失函数。"""
+
+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())

engine/metrics.py

@@ -0,0 +1,171 @@
+"""指标：分类与定位。"""
+
+from __future__ import annotations
+
+import torch
+
+
+class AverageMeter:
+    def __init__(self) -> None:
+        self.total = 0.0
+        self.count = 0
+
+    def update(self, value: float, n: int = 1) -> None:
+        self.total += float(value) * n
+        self.count += n
+
+    @property
+    def avg(self) -> float:
+        return self.total / self.count if self.count else 0.0
+
+
+def box_iou(boxes_a: torch.Tensor, boxes_b: torch.Tensor) -> torch.Tensor:
+    """(N, 4) 与 (M, 4) 的两两 IoU，返回 (N, M)。"""
+    if boxes_a.numel() == 0 or boxes_b.numel() == 0:
+        return boxes_a.new_zeros((boxes_a.shape[0], boxes_b.shape[0]))
+    area_a = (boxes_a[:, 2] - boxes_a[:, 0]).clamp(min=0) * (boxes_a[:, 3] - boxes_a[:, 1]).clamp(min=0)
+    area_b = (boxes_b[:, 2] - boxes_b[:, 0]).clamp(min=0) * (boxes_b[:, 3] - boxes_b[:, 1]).clamp(min=0)
+
+    lt = torch.max(boxes_a[:, None, :2], boxes_b[None, :, :2])
+    rb = torch.min(boxes_a[:, None, 2:], boxes_b[None, :, 2:])
+    wh = (rb - lt).clamp(min=0)
+    inter = wh[..., 0] * wh[..., 1]
+    return inter / (area_a[:, None] + area_b[None, :] - inter).clamp(min=1e-6)
+
+
+class DetectionMetric:
+    """按 IoU 阈值做贪心匹配，统计 precision / recall / F1 与平均 IoU。
+
+    这是工程上够用的近似指标；正式对比模型时建议再跑一次 COCO mAP。
+    """
+
+    def __init__(self, iou_threshold: float = 0.5) -> None:
+        self.iou_threshold = iou_threshold
+        self.reset()
+
+    def reset(self) -> None:
+        self.tp = 0
+        self.fp = 0
+        self.fn = 0
+        self.iou_sum = 0.0
+
+    def update(
+        self,
+        pred_boxes: torch.Tensor,
+        pred_labels: torch.Tensor,
+        gt_boxes: torch.Tensor,
+        gt_labels: torch.Tensor,
+    ) -> None:
+        if gt_boxes.numel() == 0:
+            self.fp += int(pred_boxes.shape[0])
+            return
+        if pred_boxes.numel() == 0:
+            self.fn += int(gt_boxes.shape[0])
+            return
+
+        ious = box_iou(pred_boxes.float().cpu(), gt_boxes.float().cpu())
+        same_class = pred_labels.cpu()[:, None] == gt_labels.cpu()[None, :]
+        ious = ious * same_class.float()
+
+        matched_gt: set[int] = set()
+        for pred_idx in range(ious.shape[0]):
+            best_iou, best_gt = ious[pred_idx].max(dim=0)
+            gt_idx = int(best_gt)
+            if float(best_iou) >= self.iou_threshold and gt_idx not in matched_gt:
+                matched_gt.add(gt_idx)
+                self.tp += 1
+                self.iou_sum += float(best_iou)
+            else:
+                self.fp += 1
+        self.fn += int(gt_boxes.shape[0]) - len(matched_gt)
+
+    def compute(self) -> dict[str, float]:
+        precision = self.tp / max(self.tp + self.fp, 1)
+        recall = self.tp / max(self.tp + self.fn, 1)
+        f1 = 2 * precision * recall / max(precision + recall, 1e-9)
+        return {
+            "precision": precision,
+            "recall": recall,
+            "f1": f1,
+            "mean_iou": self.iou_sum / max(self.tp, 1),
+        }
+
+
+class MultiLabelMetric:
+    """多类模式：逐类阈值化后统计 micro precision / recall / F1 与整图完全正确率。"""
+
+    def __init__(self, threshold: float = 0.5) -> None:
+        self.threshold = threshold
+        self.reset()
+
+    def reset(self) -> None:
+        self.tp = 0
+        self.fp = 0
+        self.fn = 0
+        self.exact = 0
+        self.total = 0
+
+    def update(self, logits: torch.Tensor, targets: torch.Tensor) -> None:
+        pred = (logits.sigmoid() >= self.threshold).float().cpu()
+        gt = targets.float().cpu()
+        self.tp += int((pred * gt).sum())
+        self.fp += int((pred * (1 - gt)).sum())
+        self.fn += int(((1 - pred) * gt).sum())
+        self.exact += int((pred == gt).all(dim=1).sum())
+        self.total += int(gt.shape[0])
+
+    def compute(self) -> dict[str, float]:
+        precision = self.tp / max(self.tp + self.fp, 1)
+        recall = self.tp / max(self.tp + self.fn, 1)
+        return {
+            "cls_precision": precision,
+            "cls_recall": recall,
+            "cls_f1": 2 * precision * recall / max(precision + recall, 1e-9),
+            "cls_exact_match": self.exact / max(self.total, 1),
+        }
+
+
+class SingleLabelMetric:
+    """单类模式：互斥分类的准确率与 macro precision / recall / F1。
+
+    指标名与 `MultiLabelMetric` 保持一致（cls_precision / cls_recall / cls_f1），
+    因此切换 label_mode 后 `train.monitor` 不需要改。
+    """
+
+    def __init__(self, num_outputs: int) -> None:
+        self.num_outputs = num_outputs
+        self.reset()
+
+    def reset(self) -> None:
+        self.tp = torch.zeros(self.num_outputs)
+        self.fp = torch.zeros(self.num_outputs)
+        self.fn = torch.zeros(self.num_outputs)
+        self.correct = 0
+        self.total = 0
+
+    def update(self, logits: torch.Tensor, targets: torch.Tensor) -> None:
+        pred = logits.argmax(dim=1).cpu()
+        gt = targets.long().cpu()
+        for cls in range(self.num_outputs):
+            pred_cls = pred == cls
+            gt_cls = gt == cls
+            self.tp[cls] += int((pred_cls & gt_cls).sum())
+            self.fp[cls] += int((pred_cls & ~gt_cls).sum())
+            self.fn[cls] += int((~pred_cls & gt_cls).sum())
+        self.correct += int((pred == gt).sum())
+        self.total += int(gt.numel())
+
+    def compute(self) -> dict[str, float]:
+        # 只对验证集中出现过的类别做 macro 平均，避免缺席类别把指标拉到 0。
+        present = (self.tp + self.fn) > 0
+        precision = self.tp / (self.tp + self.fp).clamp(min=1)
+        recall = self.tp / (self.tp + self.fn).clamp(min=1)
+        f1 = 2 * precision * recall / (precision + recall).clamp(min=1e-9)
+        if present.any():
+            precision, recall, f1 = precision[present], recall[present], f1[present]
+        return {
+            "cls_accuracy": self.correct / max(self.total, 1),
+            "cls_precision": float(precision.mean()),
+            "cls_recall": float(recall.mean()),
+            "cls_f1": float(f1.mean()),
+        }

engine/optim.py

@@ -0,0 +1,61 @@
+"""优化器与学习率调度。"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+import torch
+from torch import nn
+
+
+def build_optimizer(model: nn.Module, cfg: Mapping[str, Any]) -> torch.optim.Optimizer:
+    """`cfg` 形如 `{name: sgd, lr: 0.01, momentum: 0.9, weight_decay: 1e-4}`。
+
+    BatchNorm 与 bias 不做 weight decay——小数据集上这项影响明显。
+    """
+    cfg = dict(cfg)
+    name = str(cfg.pop("name", "adamw")).lower()
+    weight_decay = float(cfg.pop("weight_decay", 0.0))
+
+    decay, no_decay = [], []
+    for module in model.modules():
+        for param_name, param in module.named_parameters(recurse=False):
+            if not param.requires_grad:
+                continue
+            if param_name == "bias" or isinstance(module, nn.modules.batchnorm._BatchNorm):
+                no_decay.append(param)
+            else:
+                decay.append(param)
+    groups = [
+        {"params": decay, "weight_decay": weight_decay},
+        {"params": no_decay, "weight_decay": 0.0},
+    ]
+
+    if name == "sgd":
+        cfg.setdefault("momentum", 0.9)
+        cfg.setdefault("nesterov", True)
+        return torch.optim.SGD(groups, **cfg)
+    if name == "adam":
+        return torch.optim.Adam(groups, **cfg)
+    if name == "adamw":
+        return torch.optim.AdamW(groups, **cfg)
+    raise ValueError(f"未知优化器: {name}")
+
+
+def build_scheduler(optimizer: torch.optim.Optimizer, cfg: Mapping[str, Any] | None, epochs: int):
+    """返回按 epoch 步进的 scheduler；`cfg` 为空则不使用。"""
+    if not cfg:
+        return None
+    cfg = dict(cfg)
+    name = str(cfg.pop("name", "cosine")).lower()
+    if name == "cosine":
+        cfg.setdefault("T_max", epochs)
+        return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, **cfg)
+    if name == "step":
+        cfg.setdefault("step_size", max(epochs // 3, 1))
+        return torch.optim.lr_scheduler.StepLR(optimizer, **cfg)
+    if name == "multistep":
+        cfg.setdefault("milestones", [int(epochs * 0.6), int(epochs * 0.85)])
+        return torch.optim.lr_scheduler.MultiStepLR(optimizer, **cfg)
+    raise ValueError(f"未知调度器: {name}")

engine/trainer.py

@@ -0,0 +1,198 @@
+"""训练循环。
+
+注意：训练入口必须先通过 `utils.device.require_cuda()` 的 GPU 检查。
+"""
+
+from __future__ import annotations
+
+import math
+import time
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+
+from engine.checkpoint import save_checkpoint
+from engine.evaluator import Evaluator, _to_device
+from engine.losses import DetectionLoss
+from engine.metrics import AverageMeter
+from engine.optim import build_optimizer, build_scheduler
+from engine.visualizer import SwanLabLogger
+from utils.labels import MULTILABEL
+from utils.logging import get_logger
+
+logger = get_logger("engine.trainer")
+
+
+class Trainer:
+    """Args:
+    model: 已经 `.to(device)` 的模型。
+    train_loader / val_loader: 数据加载器，val 可为 None。
+    cfg: `train` 段配置（epochs / optimizer / scheduler / amp / ...）。
+    work_dir: 日志与权重输出目录。
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        train_loader: DataLoader,
+        val_loader: DataLoader | None,
+        cfg: Mapping[str, Any],
+        device: torch.device,
+        work_dir: str | Path,
+        full_config: dict | None = None,
+    ) -> None:
+        self.model = model
+        self.train_loader = train_loader
+        self.val_loader = val_loader
+        self.cfg = dict(cfg)
+        self.device = device
+        self.work_dir = Path(work_dir)
+        self.work_dir.mkdir(parents=True, exist_ok=True)
+        self.full_config = full_config
+
+        self.epochs = int(self.cfg.get("epochs", 50))
+        self.log_interval = int(self.cfg.get("log_interval", 20))
+        self.clip_grad = float(self.cfg.get("clip_grad", 0.0))
+        self.monitor = str(self.cfg.get("monitor", "f1"))
+
+        # 分类损失形式跟随模型的 label_mode（yaml 里的 label_mode 决定多类 / 单类）。
+        loss_cfg = dict(self.cfg.get("loss", {}))
+        loss_cfg.setdefault("label_mode", getattr(model, "label_mode", MULTILABEL))
+        self.criterion = DetectionLoss(**loss_cfg)
+        self.optimizer = build_optimizer(model, self.cfg.get("optimizer", {"name": "adamw", "lr": 1e-3}))
+        self.scheduler = build_scheduler(self.optimizer, self.cfg.get("scheduler"), self.epochs)
+
+        self.amp = bool(self.cfg.get("amp", True)) and device.type == "cuda"
+        self.scaler = torch.amp.GradScaler("cuda", enabled=self.amp)
+
+        self.evaluator = (
+            Evaluator(
+                model,
+                val_loader,
+                device,
+                criterion=self.criterion,
+                **dict(self.cfg.get("eval", {})),
+            )
+            if val_loader is not None
+            else None
+        )
+        self.visualizer = SwanLabLogger(
+            work_dir=self.work_dir,
+            config=full_config,
+            **dict(self.cfg.get("swanlab", {})),
+        )
+        self.best_metric = float("-inf")
+        self.start_epoch = 0
+        self.global_step = 0
+
+    def train_one_epoch(self, epoch: int) -> dict[str, float]:
+        self.model.train()
+        meters: dict[str, AverageMeter] = {}
+        start = time.time()
+
+        for step, batch in enumerate(self.train_loader):
+            batch = _to_device(batch, self.device)
+
+            with torch.amp.autocast("cuda", enabled=self.amp):
+                outputs = self.model(batch["images"])
+                losses = self.criterion(outputs, batch)
+
+            self.optimizer.zero_grad(set_to_none=True)
+            self.scaler.scale(losses["total"]).backward()
+            if self.clip_grad > 0:
+                self.scaler.unscale_(self.optimizer)
+                nn.utils.clip_grad_norm_(self.model.parameters(), self.clip_grad)
+            self.scaler.step(self.optimizer)
+            self.scaler.update()
+
+            batch_size = batch["images"].shape[0]
+            step_losses = {key: float(value.detach()) for key, value in losses.items()}
+            for key, value in step_losses.items():
+                meters.setdefault(key, AverageMeter()).update(value, batch_size)
+
+            self.global_step += 1
+            if step % self.log_interval == 0:
+                self.visualizer.log(
+                    {**step_losses, "lr": self.optimizer.param_groups[0]["lr"]},
+                    step=self.global_step,
+                    prefix="train/",
+                )
+                summary = " ".join(f"{k}={m.avg:.4f}" for k, m in meters.items())
+                logger.info(
+                    "epoch %d [%d/%d] lr=%.2e %s",
+                    epoch,
+                    step,
+                    len(self.train_loader),
+                    self.optimizer.param_groups[0]["lr"],
+                    summary,
+                )
+
+        logger.info("epoch %d 训练耗时 %.1fs", epoch, time.time() - start)
+        return {key: meter.avg for key, meter in meters.items()}
+
+    def _gpu_memory(self) -> dict[str, float]:
+        """本轮训练的显存峰值（GB）；推理端的详细分析见 scripts/profile_inference.py。"""
+        if self.device.type != "cuda":
+            return {}
+        stats = {
+            "gpu_mem_allocated_gb": torch.cuda.max_memory_allocated(self.device) / 1024**3,
+            "gpu_mem_reserved_gb": torch.cuda.max_memory_reserved(self.device) / 1024**3,
+        }
+        torch.cuda.reset_peak_memory_stats(self.device)
+        return stats
+
+    def fit(self) -> float:
+        logger.info("开始训练：共 %d 个 epoch，输出目录 %s", self.epochs, self.work_dir)
+        for epoch in range(self.start_epoch, self.epochs):
+            train_metrics = self.train_one_epoch(epoch)
+            if self.scheduler is not None:
+                self.scheduler.step()
+
+            metrics: dict[str, float] = {}
+            if self.evaluator is not None:
+                metrics = self.evaluator.evaluate()
+                logger.info(
+                    "epoch %d 验证: %s",
+                    epoch,
+                    " ".join(f"{k}={v:.4f}" for k, v in metrics.items()),
+                )
+
+            self.visualizer.log(
+                {
+                    **{f"train/epoch_{k}": v for k, v in train_metrics.items()},
+                    **{f"val/{k}": v for k, v in metrics.items()},
+                    **{f"train/{k}": v for k, v in self._gpu_memory().items()},
+                    "epoch": epoch,
+                },
+                step=self.global_step,
+            )
+
+            save_checkpoint(
+                self.work_dir / "last.pth",
+                self.model,
+                self.optimizer,
+                epoch=epoch,
+                best_metric=self.best_metric,
+                config=self.full_config,
+            )
+            score = metrics.get(self.monitor)
+            if score is not None and score > self.best_metric:
+                self.best_metric = score
+                save_checkpoint(
+                    self.work_dir / "best.pth",
+                    self.model,
+                    self.optimizer,
+                    epoch=epoch,
+                    best_metric=self.best_metric,
+                    config=self.full_config,
+                )
+                logger.info("epoch %d 刷新最佳 %s=%.4f，已保存 best.pth", epoch, self.monitor, score)
+
+        if math.isfinite(self.best_metric):
+            self.visualizer.log({f"best_{self.monitor}": self.best_metric}, step=self.global_step)
+        self.visualizer.finish()
+        return self.best_metric

engine/visualizer.py

@@ -0,0 +1,100 @@
+"""训练可视化：SwanLab 封装。
+
+配置示例（`train.swanlab`）：
+
+```yaml
+train:
+  swanlab:
+    enabled: true
+    project: film-defect
+    name: baseline          # 缺省用 work_dir 的目录名
+    mode: online            # online(云端) | local(本地看板) | offline(仅存盘) | disabled
+    tags: [tiny_conv, center-head]
+    log_dir: null           # 缺省 <work_dir>/swanlog
+```
+
+未安装 swanlab、或 `enabled: false` 时整体退化为空操作，训练流程不受影响。
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from utils.logging import get_logger
+
+logger = get_logger("engine.visualizer")
+
+VALID_MODES = ("online", "local", "offline", "disabled")
+
+
+class SwanLabLogger:
+    """把标量指标写到 SwanLab；不可用时静默降级。
+
+    Args:
+        enabled: 关掉后所有方法都是空操作。
+        work_dir: 用于推导默认的实验名与 `log_dir`。
+        config: 记录到实验里的超参（一般直接传完整 yaml 配置）。
+        其余参数透传给 `swanlab.init`。
+    """
+
+    def __init__(
+        self,
+        enabled: bool = True,
+        work_dir: str | Path = "runs",
+        config: dict | None = None,
+        project: str = "film-defect",
+        name: str | None = None,
+        mode: str = "online",
+        log_dir: str | Path | None = None,
+        **init_kwargs: Any,
+    ) -> None:
+        work_dir = Path(work_dir)
+        self.run = None
+        self.enabled = bool(enabled) and mode != "disabled"
+        if not self.enabled:
+            logger.info("SwanLab 可视化未启用")
+            return
+
+        if mode not in VALID_MODES:
+            raise ValueError(f"未知 swanlab mode: {mode}，可选 {VALID_MODES}")
+
+        try:
+            import swanlab
+        except ImportError:
+            logger.warning("未安装 swanlab，跳过可视化（uv sync 后可用；或设 train.swanlab.enabled=false）")
+            self.enabled = False
+            return
+
+        try:
+            self.run = swanlab.init(
+                project=project,
+                name=name or work_dir.name,
+                mode=mode,
+                config=config or {},
+                log_dir=str(log_dir or work_dir / "swanlog"),
+                **init_kwargs,
+            )
+            self._swanlab = swanlab
+            logger.info("SwanLab 已启动（project=%s, mode=%s）", project, mode)
+        except Exception as exc:  # 联网失败 / 未登录不应中断训练
+            logger.warning("SwanLab 初始化失败，训练继续但不记录可视化: %s", exc)
+            self.enabled = False
+
+    def log(self, metrics: dict[str, float], step: int | None = None, prefix: str = "") -> None:
+        if not self.enabled:
+            return
+        payload = {f"{prefix}{key}": float(value) for key, value in metrics.items()}
+        try:
+            self._swanlab.log(payload, step=step)
+        except Exception as exc:  # 记录失败同样不中断训练
+            logger.warning("SwanLab 记录失败: %s", exc)
+            self.enabled = False
+
+    def finish(self) -> None:
+        if not self.enabled:
+            return
+        try:
+            self._swanlab.finish()
+        except Exception as exc:
+            logger.warning("SwanLab 收尾失败: %s", exc)