Shaun ZhouandClaude Opus 5
105cb38831
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>
291 lines
12 KiB
Python
291 lines
12 KiB
Python
"""训练循环。支持 cuda / mps / cpu 三种后端。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
import time
|
||
import warnings
|
||
from collections.abc import Mapping
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import torch
|
||
from torch import nn
|
||
from torch.utils.data import DataLoader
|
||
|
||
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
|
||
from engine.optim import build_optimizer, build_scheduler
|
||
from engine.visualizer import SwanLabLogger
|
||
from utils.device import PeakMemory, resolve_amp
|
||
from utils.labels import MULTILABEL
|
||
from utils.logging import get_logger
|
||
|
||
logger = get_logger("engine.trainer")
|
||
|
||
|
||
class Trainer:
|
||
"""Args:
|
||
model: 已经 `.to(device)` 的模型。
|
||
train_loader / val_loader: 数据加载器，val 可为 None。
|
||
cfg: `train` 段配置（epochs / optimizer / scheduler / amp / ...）。
|
||
work_dir: 日志与权重输出目录。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
model: nn.Module,
|
||
train_loader: DataLoader,
|
||
val_loader: DataLoader | None,
|
||
cfg: Mapping[str, Any],
|
||
device: torch.device,
|
||
work_dir: str | Path,
|
||
full_config: dict | None = None,
|
||
) -> None:
|
||
self.model = model
|
||
self.train_loader = train_loader
|
||
self.val_loader = val_loader
|
||
self.cfg = dict(cfg)
|
||
self.device = device
|
||
self.work_dir = Path(work_dir)
|
||
self.work_dir.mkdir(parents=True, exist_ok=True)
|
||
self.full_config = full_config
|
||
|
||
self.epochs = int(self.cfg.get("epochs", 50))
|
||
self.log_interval = int(self.cfg.get("log_interval", 20))
|
||
self.clip_grad = float(self.cfg.get("clip_grad", 0.0))
|
||
self.monitor = str(self.cfg.get("monitor", "f1"))
|
||
|
||
# 分类损失形式跟随模型的 label_mode（yaml 里的 label_mode 决定多类 / 单类）。
|
||
loss_cfg = dict(self.cfg.get("loss", {}))
|
||
loss_cfg.setdefault("label_mode", getattr(model, "label_mode", MULTILABEL))
|
||
self.criterion = DetectionLoss(**loss_cfg)
|
||
self.optimizer = build_optimizer(model, self.cfg.get("optimizer", {"name": "adamw", "lr": 1e-3}))
|
||
self.scheduler = build_scheduler(self.optimizer, self.cfg.get("scheduler"), self.epochs)
|
||
|
||
# amp: auto 时随后端（CUDA 开、MPS 关——MPS 上 autocast 实测比 fp32 慢）
|
||
self.amp = resolve_amp(self.cfg.get("amp", "auto"), device)
|
||
self.scaler = torch.amp.GradScaler(device.type, enabled=self.amp)
|
||
self.peak_memory = PeakMemory(device)
|
||
|
||
self.evaluator = (
|
||
Evaluator(
|
||
model,
|
||
val_loader,
|
||
device,
|
||
criterion=self.criterion,
|
||
**dict(self.cfg.get("eval", {})),
|
||
)
|
||
if val_loader is not None
|
||
else None
|
||
)
|
||
self.visualizer = SwanLabLogger(
|
||
work_dir=self.work_dir,
|
||
config=full_config,
|
||
**dict(self.cfg.get("swanlab", {})),
|
||
)
|
||
self.best_metric = float("-inf")
|
||
self.start_epoch = 0
|
||
self.global_step = 0
|
||
|
||
def train_one_epoch(self, epoch: int) -> dict[str, float]:
|
||
self.model.train()
|
||
meters: dict[str, AverageMeter] = {}
|
||
start = time.time()
|
||
|
||
for step, batch in enumerate(self.train_loader):
|
||
batch = _to_device(batch, self.device)
|
||
|
||
with torch.amp.autocast(self.device.type, enabled=self.amp):
|
||
outputs = self.model(batch["images"])
|
||
losses = self.criterion(outputs, batch)
|
||
|
||
self.optimizer.zero_grad(set_to_none=True)
|
||
self.scaler.scale(losses["total"]).backward()
|
||
if self.clip_grad > 0:
|
||
self.scaler.unscale_(self.optimizer)
|
||
nn.utils.clip_grad_norm_(self.model.parameters(), self.clip_grad)
|
||
self.scaler.step(self.optimizer)
|
||
self.scaler.update()
|
||
|
||
batch_size = batch["images"].shape[0]
|
||
step_losses = {key: float(value.detach()) for key, value in losses.items()}
|
||
for key, value in step_losses.items():
|
||
meters.setdefault(key, AverageMeter()).update(value, batch_size)
|
||
|
||
self.global_step += 1
|
||
if step % self.log_interval == 0:
|
||
self.visualizer.log(
|
||
{**step_losses, "lr": self.optimizer.param_groups[0]["lr"]},
|
||
step=self.global_step,
|
||
prefix="train/",
|
||
)
|
||
summary = " ".join(f"{k}={m.avg:.4f}" for k, m in meters.items())
|
||
logger.info(
|
||
"epoch %d [%d/%d] lr=%.2e %s",
|
||
epoch,
|
||
step,
|
||
len(self.train_loader),
|
||
self.optimizer.param_groups[0]["lr"],
|
||
summary,
|
||
)
|
||
|
||
logger.info("epoch %d 训练耗时 %.1fs", epoch, time.time() - start)
|
||
return {key: meter.avg for key, meter in meters.items()}
|
||
|
||
def _gpu_memory(self) -> dict[str, float]:
|
||
"""本轮训练的显存峰值（GB）；推理端的详细分析见 scripts/profile_inference.py。
|
||
|
||
CUDA 读内置峰值计数器；MPS 没有峰值 API，用逐 step 采样的最大值（是下界）；
|
||
CPU 不统计。
|
||
"""
|
||
if self.device.type == "cpu":
|
||
return {}
|
||
stats = {
|
||
"gpu_mem_allocated_gb": self.peak_memory.peak_mb / 1024,
|
||
"gpu_mem_reserved_gb": self.peak_memory.reserved_mb / 1024,
|
||
}
|
||
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):
|
||
train_metrics = self.train_one_epoch(epoch)
|
||
if self.scheduler is not None:
|
||
self.scheduler.step()
|
||
|
||
metrics: dict[str, float] = {}
|
||
if self.evaluator is not None:
|
||
metrics = self.evaluator.evaluate()
|
||
logger.info(
|
||
"epoch %d 验证: %s",
|
||
epoch,
|
||
" ".join(f"{k}={v:.4f}" for k, v in metrics.items()),
|
||
)
|
||
|
||
self.visualizer.log(
|
||
{
|
||
**{f"train/epoch_{k}": v for k, v in train_metrics.items()},
|
||
**{f"val/{k}": v for k, v in metrics.items()},
|
||
**{f"train/{k}": v for k, v in self._gpu_memory().items()},
|
||
"epoch": epoch,
|
||
},
|
||
step=self.global_step,
|
||
)
|
||
|
||
save_checkpoint(
|
||
self.work_dir / "last.pth",
|
||
self.model,
|
||
self.optimizer,
|
||
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:
|
||
self.best_metric = score
|
||
save_checkpoint(
|
||
self.work_dir / "best.pth",
|
||
self.model,
|
||
self.optimizer,
|
||
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)
|
||
|
||
if math.isfinite(self.best_metric):
|
||
self.visualizer.log({f"best_{self.monitor}": self.best_metric}, step=self.global_step)
|
||
self.visualizer.finish()
|
||
return self.best_metric
|