fix(trainer): 检测训练续训以配置为准，并恢复学习率调度进度
train.py --resume 走的是与 SimCLR 不同的一套代码，同样的问题在这里更严重： - load_checkpoint 直接 optimizer.load_state_dict，把 lr/weight_decay 连同动量缓冲 一起还原成存档值：yaml 写 lr=5e-4，实际仍按检查点的 1e-3 训 - save_checkpoint 根本没保存调度器状态，续训时 scheduler.last_epoch 从 0 开始， 而 optimizer 的 lr 又来自检查点，两者互相矛盾，cosine 曲线完全错乱 改为与 SimCLR 一致的「配置文件是唯一真相源」： - save_checkpoint 补上 scheduler 状态、global_step 与当时的 train 配置快照 - 新增 Trainer.resume()：加载优化器状态后按组写回配置里的超参（只留动量缓冲， 并保住 BN/bias 那组 weight_decay=0 的结构），调度器按当前配置重建再快进到已训 epoch - 配置与检查点不一致时逐项打日志说明按当前配置执行 - train.py 改用 trainer.resume()，不再自己拼恢复逻辑 实测续训到 10 轮时 lr=3.273e-04，与 cosine 在 4/10 处的理论值精确一致； weight_decay、T_max、调度进度、动量缓冲均正确。 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
01cfff9b02
commit
105cb38831
+89
-1
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import math
|
||||
import time
|
||||
import warnings
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
@@ -12,7 +13,7 @@ import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from engine.checkpoint import save_checkpoint
|
||||
from engine.checkpoint import load_checkpoint, save_checkpoint
|
||||
from engine.evaluator import Evaluator, _to_device
|
||||
from engine.losses import DetectionLoss
|
||||
from engine.metrics import AverageMeter
|
||||
@@ -149,6 +150,87 @@ class Trainer:
|
||||
self.peak_memory.reset()
|
||||
return stats
|
||||
|
||||
def resume(self, path: str | Path) -> dict:
|
||||
"""从检查点续训。
|
||||
|
||||
**配置文件是唯一真相源**：检查点只提供"进度"——模型权重、优化器动量缓冲、
|
||||
已训到第几个 epoch、global_step、best_metric；而 lr、weight_decay、总轮数、
|
||||
调度器一律以当前 yaml 为准。否则改了配置再续训，改动会被检查点静默覆盖。
|
||||
"""
|
||||
state = load_checkpoint(path, self.model, map_location=self.device)
|
||||
|
||||
if state.get("optimizer"):
|
||||
# optimizer.load_state_dict 会整体覆盖 param_groups，把 lr/weight_decay
|
||||
# 一并还原成存档时的值。先抄下配置构造出的超参，加载完按组写回，
|
||||
# 只留下动量缓冲。按组写回也保住了 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)
|
||||
|
||||
self.start_epoch = int(state.get("epoch", -1)) + 1
|
||||
self.global_step = int(state.get("global_step", 0))
|
||||
best = state.get("best_metric")
|
||||
self.best_metric = float(best) if best is not None and math.isfinite(float(best)) else float("-inf")
|
||||
|
||||
self._restore_scheduler()
|
||||
self._log_config_changes(state)
|
||||
logger.info(
|
||||
"已从 %s 恢复：起始 epoch=%d，best_%s=%.4f",
|
||||
path,
|
||||
self.start_epoch,
|
||||
self.monitor,
|
||||
self.best_metric,
|
||||
)
|
||||
return state
|
||||
|
||||
def _restore_scheduler(self) -> None:
|
||||
"""按**当前配置**重建学习率调度，再快进到已训到的 epoch。
|
||||
|
||||
不恢复检查点里的调度器状态：它带着当时的 T_max 与 base_lrs，一旦 yaml 里改了
|
||||
总轮数或学习率就会被静默覆盖。cosine / step / multistep 都是 epoch 的确定性函数，
|
||||
快进与恢复状态等价，所以统一走重建这一条路。
|
||||
"""
|
||||
if self.scheduler is None:
|
||||
return
|
||||
# cosine 的链式实现基于当前 lr 递推，必须先复位回 initial_lr，
|
||||
# 否则从检查点带来的 lr 出发会算出一条错的曲线。
|
||||
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()
|
||||
|
||||
def _log_config_changes(self, state: dict) -> None:
|
||||
"""检查点里的配置与当前配置不一致时说清楚用的是哪一份。"""
|
||||
archived = state.get("train_cfg") or {}
|
||||
changes = []
|
||||
if archived.get("epochs") not in (None, self.epochs):
|
||||
changes.append(f"总轮数 {archived['epochs']} -> {self.epochs}")
|
||||
for key, value in (archived.get("optimizer") or {}).items():
|
||||
current = dict(self.cfg.get("optimizer", {})).get(key)
|
||||
if current is not None and current != value:
|
||||
changes.append(f"optimizer.{key} {value} -> {current}")
|
||||
archived_scheduler = (archived.get("scheduler") or {}).get("name")
|
||||
current_scheduler = (dict(self.cfg.get("scheduler") or {})).get("name")
|
||||
if archived_scheduler and 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 fit(self) -> float:
|
||||
logger.info("开始训练：共 %d 个 epoch，输出目录 %s", self.epochs, self.work_dir)
|
||||
for epoch in range(self.start_epoch, self.epochs):
|
||||
@@ -182,6 +264,9 @@ class Trainer:
|
||||
epoch=epoch,
|
||||
best_metric=self.best_metric,
|
||||
config=self.full_config,
|
||||
scheduler=self.scheduler,
|
||||
global_step=self.global_step,
|
||||
train_cfg=dict(self.cfg),
|
||||
)
|
||||
score = metrics.get(self.monitor)
|
||||
if score is not None and score > self.best_metric:
|
||||
@@ -193,6 +278,9 @@ class Trainer:
|
||||
epoch=epoch,
|
||||
best_metric=self.best_metric,
|
||||
config=self.full_config,
|
||||
scheduler=self.scheduler,
|
||||
global_step=self.global_step,
|
||||
train_cfg=dict(self.cfg),
|
||||
)
|
||||
logger.info("epoch %d 刷新最佳 %s=%.4f，已保存 best.pth", epoch, self.monitor, score)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user