fix(simclr): 换优化器种类时不再加载不兼容的动量缓冲

adamw 换成 sgd 时 load_state_dict 不会报错,SGD 会收下 AdamW 的 exp_avg/exp_avg_sq
却永远用不上——动量实际从零开始,但状态里白占内存,之后存盘还会一路带下去。
改为检测到优化器种类变化就跳过加载并给出日志,模型权重与训练进度不受影响。

README 补一张表,说明续训时哪些配置按 yaml、哪些来自存档、哪些不能改。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
2026-09-05 11:10:04 +08:00
co-authored by Claude Opus 5
parent b7505a743d
commit 004d7a1ce7
3 changed files with 57 additions and 7 deletions
+18 -7
View File
@@ -224,13 +224,24 @@ class SimCLRTrainer:
# 一并还原成存档时的值。先把配置构造出来的超参抄下来,加载完再按组写回去,
# 这样只留下动量缓冲(真正的"进度"),超参仍以 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)
archived_optimizer = (state.get("optimizer_cfg") or {}).get("name")
current_optimizer = dict(self.cfg.get("optimizer", {})).get("name")
if archived_optimizer and current_optimizer and archived_optimizer != current_optimizer:
# 换了优化器种类,动量缓冲的含义完全不同(AdamW 的 exp_avg 对 SGD 毫无意义),
# 硬塞进去只会白占内存。直接不加载,从零开始积累。
logger.warning(
"优化器由 %s 换成了 %s,动量缓冲无法迁移、将从零积累(模型权重与训练进度不受影响)",
archived_optimizer,
current_optimizer,
)
else:
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"])