fix(simclr): 续训时改变总轮数不再让学习率卡在 0
resume 会恢复调度器状态,其中包含存档时的 T_max。原计划 10 轮的 cosine 跑完后 学习率已降到 0,此时把 epochs 改成 30 接着练,学习率仍是 0——训练照常进行、 日志照常打印,但一步都学不动。是静默失败,比报错更难发现。 - 总轮数与存档一致时行为不变,直接恢复调度器状态 - 不一致时按新的总轮数重建调度器,并快进到当前 epoch;重建前先把 lr 复位回 initial_lr,因为 cosine 的链式实现基于当前 lr 递推,从 0 出发只会一路留在 0 - 实测续训到 30 轮时 lr 恢复为 7.50e-04,与 cosine 在 10/30 处的理论值一致 - 补三个回归测试:改变总轮数后的学习率、总轮数不变时直接恢复状态、 save_interval 落的中间盘可直接续训 - 修正一个旧测试:它断言续训后 lr 仍等于旧计划末尾的 0,编码的是错误行为 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
6b11128357
commit
865cdd7210
+35
-2
@@ -7,6 +7,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import warnings
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
@@ -210,13 +211,12 @@ class SimCLRTrainer:
|
||||
self.backbone.load_state_dict(state["backbone"])
|
||||
self.projector.load_state_dict(state["projector"])
|
||||
self.optimizer.load_state_dict(state["optimizer"])
|
||||
if self.scheduler is not None and state.get("scheduler"):
|
||||
self.scheduler.load_state_dict(state["scheduler"])
|
||||
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)
|
||||
logger.info(
|
||||
"已载入训练状态 %s:从 epoch %d 接着练到 %d", path, self.start_epoch, self.epochs
|
||||
)
|
||||
@@ -228,6 +228,39 @@ class SimCLRTrainer:
|
||||
self.epochs,
|
||||
)
|
||||
|
||||
def _restore_scheduler(self, state: dict) -> None:
|
||||
"""恢复学习率调度。
|
||||
|
||||
总轮数没变就直接恢复状态;**变了就必须按新的总轮数重建**,只把它快进到当前 epoch。
|
||||
否则会沿用存档里的 T_max:原计划 10 轮的 cosine 到第 10 轮学习率已经是 0,
|
||||
这时把 epochs 改成 30 接着练,学习率仍然是 0——训练照跑,但一步都学不动。
|
||||
"""
|
||||
if self.scheduler is None or not state.get("scheduler"):
|
||||
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。
|
||||
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)
|
||||
with warnings.catch_warnings():
|
||||
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,
|
||||
self.optimizer.param_groups[0]["lr"],
|
||||
)
|
||||
|
||||
def _checkpoint(self, epoch: int) -> None:
|
||||
"""同时落两份:给下游用的骨架,和给续训用的完整状态。"""
|
||||
self.save_backbone(epoch)
|
||||
|
||||
Reference in New Issue
Block a user