Shaun ZhouandClaude Opus 5
ef64e5bf77
- trainer: AMP、梯度裁剪、cosine/step 调度、last/best 检查点，按 monitor 指标挑选 best - losses: penalty-reduced focal loss（heatmap）+ 掩码 L1（wh/offset）+ 分类损失 （多类 BCE / 单类 CE，随 label_mode 切换） - metrics: 定位 IoU 匹配下的 P/R/F1 与 mean IoU；分类多标签 micro-F1 与单标签 accuracy， 两种模式都输出 cls_f1，故 train.monitor 无需跟着改 - optim: BatchNorm 与 bias 不做 weight decay - visualizer: SwanLab 记录损失/学习率/验证指标/显存峰值，未安装或离线时降级为警告 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
199 lines
7.3 KiB
Python
199 lines
7.3 KiB
Python
"""训练循环。
|
||
|
||
注意：训练入口必须先通过 `utils.device.require_cuda()` 的 GPU 检查。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
import time
|
||
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 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.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)
|
||
|
||
self.amp = bool(self.cfg.get("amp", True)) and device.type == "cuda"
|
||
self.scaler = torch.amp.GradScaler("cuda", enabled=self.amp)
|
||
|
||
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("cuda", 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。"""
|
||
if self.device.type != "cuda":
|
||
return {}
|
||
stats = {
|
||
"gpu_mem_allocated_gb": torch.cuda.max_memory_allocated(self.device) / 1024**3,
|
||
"gpu_mem_reserved_gb": torch.cuda.max_memory_reserved(self.device) / 1024**3,
|
||
}
|
||
torch.cuda.reset_peak_memory_stats(self.device)
|
||
return stats
|
||
|
||
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,
|
||
)
|
||
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,
|
||
)
|
||
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
|