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:
co-authored by
Claude Opus 5
parent
7d5b6ec423
commit
be6173c746
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user