fix(simclr): 续训以配置为准，不再被存档里的超参静默覆盖
optimizer.load_state_dict 会整体覆盖 param_groups，把 lr / weight_decay 连同动量缓冲 一起还原成存档时的值。改了 yaml 再续训，改动会被静默吃掉：配置写 lr=5e-4， 实际仍按存档的 1e-3 训，scheduler.base_lrs 也是旧值。 改为「配置文件是唯一真相源，存档只提供进度」： - 加载优化器状态前先抄下配置构造出的每组超参，加载后按组写回，只留下动量缓冲； 按组写回也保住了 build_optimizer 的分组结构（BN 与 bias 那组 weight_decay=0） - 调度器统一按当前配置重建再快进到已训 epoch，不再恢复存档状态（它带着旧的 T_max 与 base_lrs）。cosine/step/multistep 都是 epoch 的确定性函数， 快进与恢复状态等价，因此不存在第二种行为 - 存档里额外记下当时的 optimizer/scheduler 配置，续训时逐项列出差异， 明确说明按当前配置执行 - 补回归测试：改 lr/weight_decay/总轮数后续训，三者都按 yaml 生效， 分组结构与动量缓冲不受影响 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
703ba7076b
commit
b7505a743d
+57
-19
@@ -178,6 +178,8 @@ class SimCLRTrainer:
|
||||
|
||||
比 `simclr_backbone.pth` 多了投影头、优化器、调度器与 scaler——少了这些，
|
||||
续训等于把动量和学习率进度清零，损失会先反弹一段。
|
||||
|
||||
同时记下当时的优化器/调度器配置，续训时若与新配置不同会打日志说明按新的来。
|
||||
"""
|
||||
path = self.work_dir / "simclr_last.pth"
|
||||
torch.save(
|
||||
@@ -190,13 +192,20 @@ class SimCLRTrainer:
|
||||
"epoch": epoch,
|
||||
"global_step": self.global_step,
|
||||
"epochs_planned": self.epochs,
|
||||
"optimizer_cfg": dict(self.cfg.get("optimizer", {})),
|
||||
"scheduler_cfg": dict(self.cfg.get("scheduler") or {}),
|
||||
},
|
||||
path,
|
||||
)
|
||||
return path
|
||||
|
||||
def load_state(self, path: str | Path) -> None:
|
||||
"""载入 `save_state` 写出的完整状态，从下一个 epoch 接着练。"""
|
||||
"""载入 `save_state` 写出的完整状态，从下一个 epoch 接着练。
|
||||
|
||||
**配置文件是唯一真相源**：存档只提供"进度"——模型权重、优化器动量缓冲、
|
||||
已训到第几个 epoch、global_step；而学习率、weight_decay、总轮数、调度器
|
||||
一律以当前 yaml 为准。否则改了配置再续训，改动会被存档静默覆盖掉。
|
||||
"""
|
||||
path = Path(path)
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"续训状态文件不存在: {path}")
|
||||
@@ -210,13 +219,26 @@ class SimCLRTrainer:
|
||||
|
||||
self.backbone.load_state_dict(state["backbone"])
|
||||
self.projector.load_state_dict(state["projector"])
|
||||
|
||||
# optimizer.load_state_dict 会整体覆盖 param_groups，把 lr/weight_decay 等超参
|
||||
# 一并还原成存档时的值。先把配置构造出来的超参抄下来，加载完再按组写回去，
|
||||
# 这样只留下动量缓冲（真正的"进度"），超参仍以 yaml 为准。
|
||||
# 按组写回也保住了 build_optimizer 的分组结构（BN 与 bias 那组 weight_decay=0）。
|
||||
intended = [
|
||||
{key: value for key, value in group.items() if key != "params"}
|
||||
for group in self.optimizer.param_groups
|
||||
]
|
||||
self.optimizer.load_state_dict(state["optimizer"])
|
||||
for group, wanted in zip(self.optimizer.param_groups, intended, strict=True):
|
||||
group.update(wanted)
|
||||
|
||||
if state.get("scaler"):
|
||||
self.scaler.load_state_dict(state["scaler"])
|
||||
|
||||
self.start_epoch = int(state["epoch"]) + 1
|
||||
self.global_step = int(state.get("global_step", 0))
|
||||
self._restore_scheduler(state)
|
||||
self._log_config_changes(state)
|
||||
logger.info(
|
||||
"已载入训练状态 %s：从 epoch %d 接着练到 %d", path, self.start_epoch, self.epochs
|
||||
)
|
||||
@@ -229,22 +251,21 @@ class SimCLRTrainer:
|
||||
)
|
||||
|
||||
def _restore_scheduler(self, state: dict) -> None:
|
||||
"""恢复学习率调度。
|
||||
"""按**当前配置**重建学习率调度，再快进到已训到的 epoch。
|
||||
|
||||
总轮数没变就直接恢复状态；**变了就必须按新的总轮数重建**，只把它快进到当前 epoch。
|
||||
否则会沿用存档里的 T_max：原计划 10 轮的 cosine 到第 10 轮学习率已经是 0，
|
||||
这时把 epochs 改成 30 接着练，学习率仍然是 0——训练照跑，但一步都学不动。
|
||||
不恢复存档里的调度器状态：它带着当时的 T_max 和 base_lrs，一旦 yaml 里改了
|
||||
总轮数或学习率就会被静默覆盖。原计划 10 轮的 cosine 跑完后 lr 已是 0，
|
||||
这时把 epochs 改成 30 接着练，沿用旧 T_max 的话 lr 仍是 0——一步都学不动。
|
||||
|
||||
cosine / step / multistep 都是 epoch 的确定性函数，快进与恢复状态等价，
|
||||
所以统一走重建这一条路，没有第二种行为。
|
||||
"""
|
||||
if self.scheduler is None or not state.get("scheduler"):
|
||||
del state # 存档里的 scheduler 状态刻意不用
|
||||
if self.scheduler is None:
|
||||
return
|
||||
|
||||
planned = int(state.get("epochs_planned", self.epochs))
|
||||
if planned == self.epochs:
|
||||
self.scheduler.load_state_dict(state["scheduler"])
|
||||
return
|
||||
|
||||
# 先把 lr 复位回 initial_lr 再重建：cosine 的链式实现是基于当前 lr 递推的，
|
||||
# 而 optimizer 状态里恢复出来的 lr 可能已经退到 0，直接快进只会一路留在 0。
|
||||
# cosine 的链式实现基于当前 lr 递推，必须先复位回 initial_lr，
|
||||
# 否则从 0 出发只会一路留在 0。
|
||||
for group in self.optimizer.param_groups:
|
||||
group["lr"] = group.get("initial_lr", group["lr"])
|
||||
self.scheduler = build_scheduler(self.optimizer, self.cfg.get("scheduler"), self.epochs)
|
||||
@@ -252,13 +273,30 @@ class SimCLRTrainer:
|
||||
warnings.simplefilter("ignore") # 快进阶段的 "step before optimizer.step()" 属预期
|
||||
for _ in range(self.start_epoch):
|
||||
self.scheduler.step()
|
||||
logger.warning(
|
||||
"总轮数由 %d 改成了 %d，学习率调度按新的总轮数重建并快进到 epoch %d"
|
||||
"（不沿用存档里的 T_max），当前 lr=%.2e",
|
||||
planned,
|
||||
self.epochs,
|
||||
self.start_epoch,
|
||||
|
||||
def _log_config_changes(self, state: dict) -> None:
|
||||
"""存档配置与当前配置不一致时说清楚用的是哪一份，避免"改了没生效"的困惑。"""
|
||||
changes = []
|
||||
planned = int(state.get("epochs_planned", self.epochs))
|
||||
if planned != self.epochs:
|
||||
changes.append(f"总轮数 {planned} -> {self.epochs}")
|
||||
for key, archived in (state.get("optimizer_cfg") or {}).items():
|
||||
current = dict(self.cfg.get("optimizer", {})).get(key)
|
||||
if current is not None and current != archived:
|
||||
changes.append(f"optimizer.{key} {archived} -> {current}")
|
||||
archived_scheduler = (state.get("scheduler_cfg") or {}).get("name")
|
||||
current_scheduler = (dict(self.cfg.get("scheduler") or {})).get("name")
|
||||
if archived_scheduler != current_scheduler:
|
||||
changes.append(f"scheduler {archived_scheduler} -> {current_scheduler}")
|
||||
|
||||
if changes:
|
||||
logger.warning("配置与存档不一致，按**当前配置**执行：%s", "；".join(changes))
|
||||
logger.info(
|
||||
"续训超参：lr=%.2e（initial_lr=%.2e）weight_decay=%s，共 %d 轮",
|
||||
self.optimizer.param_groups[0]["lr"],
|
||||
self.optimizer.param_groups[0].get("initial_lr", float("nan")),
|
||||
[g["weight_decay"] for g in self.optimizer.param_groups],
|
||||
self.epochs,
|
||||
)
|
||||
|
||||
def _checkpoint(self, epoch: int) -> None:
|
||||
|
||||
Reference in New Issue
Block a user