diff --git a/engine/__init__.py b/engine/__init__.py new file mode 100644 index 0000000..8d7d7a2 --- /dev/null +++ b/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", +] diff --git a/engine/checkpoint.py b/engine/checkpoint.py new file mode 100644 index 0000000..87c9334 --- /dev/null +++ b/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 diff --git a/engine/evaluator.py b/engine/evaluator.py new file mode 100644 index 0000000..f3e8a4f --- /dev/null +++ b/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 diff --git a/engine/losses.py b/engine/losses.py new file mode 100644 index 0000000..a58b8a3 --- /dev/null +++ b/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()) diff --git a/engine/metrics.py b/engine/metrics.py new file mode 100644 index 0000000..56d3b9f --- /dev/null +++ b/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()), + } diff --git a/engine/optim.py b/engine/optim.py new file mode 100644 index 0000000..4802a5c --- /dev/null +++ b/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}") diff --git a/engine/trainer.py b/engine/trainer.py new file mode 100644 index 0000000..22e9e1a --- /dev/null +++ b/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 diff --git a/engine/visualizer.py b/engine/visualizer.py new file mode 100644 index 0000000..deaa742 --- /dev/null +++ b/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 # 缺省 /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)