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>
101 lines
3.2 KiB
Python
101 lines
3.2 KiB
Python
"""训练可视化：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)
|