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
+80
-2
@@ -239,7 +239,6 @@ def test_resume_continues_from_the_saved_epoch(tiny_image_dir, tmp_path):
|
||||
first = _make_trainer(tiny_image_dir, work_dir, epochs=2)
|
||||
first.fit()
|
||||
steps_after_2 = first.global_step
|
||||
lr_after_2 = first.optimizer.param_groups[0]["lr"]
|
||||
|
||||
resumed = SimCLRTrainer(
|
||||
BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
|
||||
@@ -253,7 +252,8 @@ def test_resume_continues_from_the_saved_epoch(tiny_image_dir, tmp_path):
|
||||
)
|
||||
assert resumed.start_epoch == 2 # 从第 3 轮开始
|
||||
assert resumed.global_step == steps_after_2 # step 计数接上
|
||||
assert resumed.optimizer.param_groups[0]["lr"] == pytest.approx(lr_after_2) # 学习率进度接上
|
||||
# 总轮数从 2 改成了 5,学习率按新计划重算(若沿用旧的 T_max,此处会是 0)
|
||||
assert resumed.optimizer.param_groups[0]["lr"] > 1e-5
|
||||
|
||||
resumed.fit()
|
||||
state = torch.load(work_dir / "simclr_last.pth", map_location="cpu", weights_only=False)
|
||||
@@ -261,6 +261,84 @@ def test_resume_continues_from_the_saved_epoch(tiny_image_dir, tmp_path):
|
||||
assert resumed.global_step > steps_after_2
|
||||
|
||||
|
||||
def test_resume_with_a_longer_schedule_rebuilds_the_lr_curve(tiny_image_dir, tmp_path):
|
||||
"""续训时调大 epochs,学习率必须按新的总轮数重算。
|
||||
|
||||
回归测试:此前直接恢复调度器状态,会沿用存档里的 T_max。原计划 10 轮的 cosine
|
||||
跑完后 lr 已经是 0,这时把 epochs 改成 30 接着练,lr 仍是 0——训练照跑但一步都学不动。
|
||||
"""
|
||||
import math
|
||||
|
||||
work_dir = tmp_path / "run"
|
||||
first = _make_trainer(tiny_image_dir, work_dir, epochs=10)
|
||||
first.fit()
|
||||
assert first.optimizer.param_groups[0]["lr"] == pytest.approx(0.0, abs=1e-9) # cosine 已走完
|
||||
|
||||
resumed = SimCLRTrainer(
|
||||
BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
|
||||
first.dataloader,
|
||||
{"epochs": 30, "amp": False, "log_interval": 999, "save_interval": 1,
|
||||
"scheduler": {"name": "cosine"}, "optimizer": {"name": "adamw", "lr": 1e-3},
|
||||
"projection": {"hidden_channels": 16, "out_channels": 8},
|
||||
"swanlab": {"enabled": False}},
|
||||
torch.device("cpu"),
|
||||
work_dir,
|
||||
resume=work_dir / "simclr_last.pth",
|
||||
)
|
||||
|
||||
assert resumed.scheduler.T_max == 30
|
||||
lr = resumed.optimizer.param_groups[0]["lr"]
|
||||
assert lr > 1e-5, "续训后学习率仍是 0,等于白跑"
|
||||
assert lr == pytest.approx(1e-3 * (1 + math.cos(math.pi * 10 / 30)) / 2, rel=1e-3)
|
||||
|
||||
|
||||
def test_resume_with_same_schedule_restores_state(tiny_image_dir, tmp_path):
|
||||
"""总轮数不变时应当直接恢复调度器状态,不重建。"""
|
||||
work_dir = tmp_path / "run"
|
||||
_make_trainer(tiny_image_dir, work_dir, epochs=6).fit()
|
||||
|
||||
resumed = SimCLRTrainer(
|
||||
BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
|
||||
_make_trainer(tiny_image_dir, work_dir, epochs=1).dataloader,
|
||||
{"epochs": 6, "amp": False, "log_interval": 999, "save_interval": 1,
|
||||
"scheduler": {"name": "cosine"},
|
||||
"projection": {"hidden_channels": 16, "out_channels": 8},
|
||||
"swanlab": {"enabled": False}},
|
||||
torch.device("cpu"),
|
||||
work_dir,
|
||||
resume=work_dir / "simclr_last.pth",
|
||||
)
|
||||
assert resumed.scheduler.T_max == 6
|
||||
assert resumed.scheduler.last_epoch == 6
|
||||
|
||||
|
||||
def test_intermediate_checkpoint_is_resumable(tiny_image_dir, tmp_path):
|
||||
"""save_interval 落的中间盘必须能直接拿来续训(不必等训练正常结束)。"""
|
||||
work_dir = tmp_path / "run"
|
||||
trainer = _make_trainer(tiny_image_dir, work_dir, epochs=10)
|
||||
# 只手动跑到第 4 轮就落盘,模拟被 kill 掉
|
||||
for epoch in range(4):
|
||||
trainer.train_one_epoch(epoch)
|
||||
trainer.scheduler.step()
|
||||
trainer._checkpoint(3)
|
||||
|
||||
resumed = SimCLRTrainer(
|
||||
BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
|
||||
trainer.dataloader,
|
||||
{"epochs": 10, "amp": False, "log_interval": 999, "save_interval": 1,
|
||||
"scheduler": {"name": "cosine"},
|
||||
"projection": {"hidden_channels": 16, "out_channels": 8},
|
||||
"swanlab": {"enabled": False}},
|
||||
torch.device("cpu"),
|
||||
work_dir,
|
||||
resume=work_dir / "simclr_last.pth",
|
||||
)
|
||||
assert resumed.start_epoch == 4
|
||||
resumed.fit()
|
||||
state = torch.load(work_dir / "simclr_last.pth", map_location="cpu", weights_only=False)
|
||||
assert state["epoch"] == 9
|
||||
|
||||
|
||||
def test_resume_rejects_backbone_only_file(tiny_image_dir, tmp_path):
|
||||
"""simclr_backbone.pth 只有骨架,不能当续训状态用,要给出可操作的报错。"""
|
||||
work_dir = tmp_path / "run"
|
||||
|
||||
Reference in New Issue
Block a user