Shaun ZhouandClaude Opus 5
105cb38831
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>
58 lines
1.6 KiB
Python
58 lines
1.6 KiB
Python
"""检查点存取。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
|
|
def save_checkpoint(
|
|
path: str | Path,
|
|
model: nn.Module,
|
|
optimizer: torch.optim.Optimizer | None = None,
|
|
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,
|
|
)
|
|
return path
|
|
|
|
|
|
def load_checkpoint(
|
|
path: str | Path,
|
|
model: nn.Module,
|
|
optimizer: torch.optim.Optimizer | None = None,
|
|
map_location: Any = "cpu",
|
|
strict: bool = True,
|
|
) -> dict:
|
|
state = torch.load(Path(path), map_location=map_location, weights_only=False)
|
|
model.load_state_dict(state["model"], strict=strict)
|
|
if optimizer is not None and state.get("optimizer"):
|
|
optimizer.load_state_dict(state["optimizer"])
|
|
return state
|