feat(simclr): 支持断点续训,恢复优化器与调度器状态

此前 simclr.checkpoint 的语义是"有权重就跳过训练",没有真正的续训:
中途停掉再启动只能从头再来,或者放弃已训的轮次。

- 新增 simclr_last.pth:骨架 + 投影头 + 优化器 + 调度器 + scaler + epoch + step,
  与给下游用的 simclr_backbone.pth 并存,两者在同一时机落盘
- 新增 simclr.resume:载入完整状态从下一个 epoch 接着练。只恢复权重不恢复动量与
  cosine 进度的话,续训时损失会先反弹一段
- resume 与 checkpoint 互斥(同时给直接报错);拿只含骨架的 simclr_backbone.pth
  去 resume 会给出可操作的报错,提示改用 checkpoint
- 续训轮数已达计划值时不再训练,只打警告

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
2026-09-03 09:34:55 +08:00
co-authored by Claude Opus 5
parent be6173c746
commit 345dee50da
6 changed files with 254 additions and 13 deletions
+10 -1
View File
@@ -5,6 +5,7 @@
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml
# 改簇数 / 换成在本地数据上自监督训练的特征(Mac 上自动走 MPS)
# 断点续训:--set method=simclr simclr.resume=runs/simclr/simclr_last.pth
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
--set num_clusters=6 method=simclr
@@ -76,7 +77,13 @@ def build_encoder(cfg, device, work_dir):
stages = simclr_cfg.pop("feature_stages", None)
backbone = BACKBONES.build(backbone_cfg).to(device)
resume = simclr_cfg.pop("resume", None)
checkpoint = simclr_cfg.pop("checkpoint", None)
if resume and checkpoint:
raise ValueError(
"simclr.resume 与 simclr.checkpoint 只能二选一:"
"resume=接着训练,checkpoint=跳过训练直接聚类。"
)
if checkpoint and Path(checkpoint).exists():
state = torch.load(checkpoint, map_location=device, weights_only=False)
backbone.load_state_dict(state["model"])
@@ -95,7 +102,9 @@ def build_encoder(cfg, device, work_dir):
collate_fn=unlabeled_collate,
)
full_config = cfg.to_dict() if hasattr(cfg, "to_dict") else dict(cfg)
SimCLRTrainer(backbone, loader, simclr_cfg, device, work_dir, full_config=full_config).fit()
SimCLRTrainer(
backbone, loader, simclr_cfg, device, work_dir, full_config=full_config, resume=resume
).fit()
return StageFeatureExtractor(backbone, stages).to(device).eval()