fix(trainer): 检测训练续训以配置为准,并恢复学习率调度进度

train.py --resume 走的是与 SimCLR 不同的一套代码,同样的问题在这里更严重:

- load_checkpoint 直接 optimizer.load_state_dict,把 lr/weight_decay 连同动量缓冲
  一起还原成存档值:yaml 写 lr=5e-4,实际仍按检查点的 1e-3 训
- save_checkpoint 根本没保存调度器状态,续训时 scheduler.last_epoch 从 0 开始,
  而 optimizer 的 lr 又来自检查点,两者互相矛盾,cosine 曲线完全错乱

改为与 SimCLR 一致的「配置文件是唯一真相源」:

- save_checkpoint 补上 scheduler 状态、global_step 与当时的 train 配置快照
- 新增 Trainer.resume():加载优化器状态后按组写回配置里的超参(只留动量缓冲,
  并保住 BN/bias 那组 weight_decay=0 的结构),调度器按当前配置重建再快进到已训 epoch
- 配置与检查点不一致时逐项打日志说明按当前配置执行
- train.py 改用 trainer.resume(),不再自己拼恢复逻辑

实测续训到 10 轮时 lr=3.273e-04,与 cosine 在 4/10 处的理论值精确一致;
weight_decay、T_max、调度进度、动量缓冲均正确。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
2026-09-07 19:06:37 +08:00
co-authored by Claude Opus 5
parent 01cfff9b02
commit 105cb38831
5 changed files with 179 additions and 7 deletions
+11
View File
@@ -16,16 +16,27 @@ def save_checkpoint(
epoch: int = 0,
best_metric: float | None = None,
config: dict | None = None,
scheduler: Any = None,
global_step: int = 0,
train_cfg: dict | None = None,
) -> Path:
"""保存检查点。
除模型与优化器外还记下调度器状态、global_step 和当时的 train 配置:
少了这些,续训时学习率调度会从头开始,而且无法判断 yaml 里的超参有没有被改过。
"""
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
torch.save(
{
"model": model.state_dict(),
"optimizer": optimizer.state_dict() if optimizer else None,
"scheduler": scheduler.state_dict() if scheduler else None,
"epoch": epoch,
"global_step": global_step,
"best_metric": best_metric,
"config": config,
"train_cfg": train_cfg,
},
path,
)