feat(simclr): 接入 SwanLab 可视化,并补上本地看板缺失的依赖

- SimCLRTrainer 记录 simclr/loss、loss_avg、lr(按 step)与 epoch_loss、
  epoch_seconds、gpu_mem_peak_mb(按 epoch),超参为完整 yaml 配置;
  正常结束与 KeyboardInterrupt 两条路径都会 finish()
- 配置新增 simclr.swanlab(字段同 train.swanlab),默认 mode: local,
  不需要登录,日志写到 <work_dir>/swanlog
- 依赖改为 swanlab[dashboard]:此前 `swanlab watch` 在干净环境里会报
  "Please install the swanboard package",README 里检测训练那节的看板命令
  同样受影响,现已随 uv sync 装好
- README 补齐本地看板/云端/离线补传三种用法与端口、局域网参数

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
2026-09-03 09:24:39 +08:00
co-authored by Claude Opus 5
parent 7d5b6ec423
commit be6173c746
7 changed files with 304 additions and 14 deletions
+55
View File
@@ -104,6 +104,61 @@ def test_simclr_trainer_runs_and_saves_backbone_only(tiny_image_dir, tmp_path):
backbone.load_state_dict(state["model"]) # 能被干净地加载回去
def test_simclr_logs_metrics_to_swanlab(tiny_image_dir, tmp_path, monkeypatch):
"""SwanLab 接线:按 step 记 loss/lr,按 epoch 记平均 loss、耗时与显存峰值,结束时收尾。"""
from torch.utils.data import DataLoader
import engine.simclr as simclr_mod
from datasets.unlabeled import UnlabeledImageDataset, unlabeled_collate
records: list[tuple[dict, int | None, str]] = []
finished: list[bool] = []
class StubLogger:
def __init__(self, **kwargs):
self.kwargs = kwargs
def log(self, metrics, step=None, prefix=""):
records.append((metrics, step, prefix))
def finish(self):
finished.append(True)
monkeypatch.setattr(simclr_mod, "SwanLabLogger", StubLogger)
dataset = UnlabeledImageDataset(
tiny_image_dir,
transforms=[{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}],
views=2,
)
loader = DataLoader(dataset, batch_size=4, shuffle=True, drop_last=True, collate_fn=unlabeled_collate)
backbone = BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]})
trainer = simclr_mod.SimCLRTrainer(
backbone,
loader,
{"epochs": 1, "amp": False, "log_interval": 1, "save_interval": 0,
"projection": {"hidden_channels": 16, "out_channels": 8},
"swanlab": {"enabled": True, "mode": "disabled"}},
torch.device("cpu"),
tmp_path / "run",
full_config={"method": "simclr"},
)
trainer.fit()
assert records, "没有任何指标被记录"
assert all(prefix == "simclr/" for _, _, prefix in records)
step_keys = {k for metrics, _, _ in records for k in metrics}
assert {"loss", "loss_avg", "lr"} <= step_keys
assert {"epoch", "epoch_loss", "epoch_seconds", "gpu_mem_peak_mb"} <= step_keys
assert all(step is not None and step > 0 for _, step, _ in records)
assert finished == [True] # fit() 结束时调用了 finish()
# 完整 yaml 配置作为超参传给了 SwanLab
assert trainer.visualizer.kwargs["config"] == {"method": "simclr"}
assert trainer.visualizer.kwargs["mode"] == "disabled"
def test_simclr_saves_intermediate_checkpoints(tiny_image_dir, tmp_path):
"""中途落盘:跑到一半被打断也要能拿到权重(此前只在全部 epoch 跑完后才存)。"""
from torch.utils.data import DataLoader