feat(engine): 训练评估循环、CenterNet 损失与指标
- 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>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
af22dc639c
commit
ef64e5bf77
@@ -0,0 +1,198 @@
|
||||
"""训练循环。
|
||||
|
||||
注意：训练入口必须先通过 `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
|
||||
Reference in New Issue
Block a user