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
+26 -2
View File
@@ -223,7 +223,9 @@ uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
```
runs/simclr/
├── simclr_backbone.pth # 每 save_interval 个 epoch 覆盖一次;Ctrl-C 也会先落盘
├── simclr_backbone.pth # 只含骨架,给下游用(聚类 / 初始化 DefectNet)
├── simclr_last.pth # 完整训练状态,给断点续训用
├── swanlog/ # SwanLab 日志
├── cluster.log
├── assignments.csv
├── summary.json
@@ -241,6 +243,26 @@ nohup uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.y
tail -f runs/simclr.log
```
#### 断点续训
训练中断(Ctrl-C、kill、机器重启)后接着练,用 `simclr.resume` 指向 `simclr_last.pth`:
```bash
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
--work-dir runs/simclr \
--set method=simclr simclr.epochs=40 loader.num_workers=8 \
simclr.resume=runs/simclr/simclr_last.pth
```
日志会打印"已载入训练状态 …:从 epoch N 接着练到 40"。恢复的不只是权重,还有**投影头、
优化器动量、学习率调度进度和 step 计数**——少了这些,续训等于把动量和 cosine 进度清零,
损失会先反弹一段。
> `resume` 与 `checkpoint` 语义不同,二选一(同时给会报错):
> **`resume`** = 载入完整状态**接着训练**(要 `simclr_last.pth`);
> **`checkpoint`** = 有权重就**跳过训练**直接聚类(要 `simclr_backbone.pth`)。
> 拿 `simclr_backbone.pth` 去 `resume` 会报错,因为它不含投影头和优化器状态。
#### 第二步:复用权重重新聚类
想换簇数或改聚类参数时**不必重训**——把权重路径填进 `simclr.checkpoint` 就会跳过训练
@@ -285,7 +307,9 @@ simclr:
temperature: 0.2 # 越小越关注难负样本
amp: auto # CUDA 上开、MPS 上关
save_interval: 5 # 每 5 个 epoch 落一次盘
checkpoint: null # 填 runs/simclr/simclr_backbone.pth 即可跳过重复训练
# 下面两个二选一,语义不同:
checkpoint: null # 有权重就跳过训练直接聚类,例 runs/simclr/simclr_backbone.pth
resume: null # 载入完整状态接着训练,例 runs/simclr/simclr_last.pth
projection: {hidden_channels: 256, out_channels: 128}
optimizer: {name: adamw, lr: 0.001, weight_decay: 0.0001}
scheduler: {name: cosine}