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