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:
co-authored by
Claude Opus 5
parent
b7505a743d
commit
004d7a1ce7
+18
-7
@@ -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"])
|
||||
|
||||
Reference in New Issue
Block a user