"""训练循环。 注意：训练入口必须先通过 `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