Shaun ZhouandClaude Opus 5
141cb907f2
移除"训练必须有 Nvidia GPU"的硬性守卫，改为按后端能力分派。
- utils/device: resolve_device("auto") 按 cuda > mps > cpu 挑最快的可用后端；
显式指定不可用的后端时直接报错，不会悄悄退回 CPU——训练慢几十倍却不吭声是最坏的情况
- AMP 随后端取默认值：CUDA 开、MPS 关。MPS 上 autocast 数值正常但实测更慢
（tiny_conv @128²：fp32 76 ms/步、fp16 90 ms、bf16 108 ms），需要时 amp: true 可强开
- PeakMemory 统一显存峰值统计：CUDA 读内置计数器，MPS 无峰值 API 改用逐 step 采样
（是下界，日志与产出里已标注）
- profile_inference 支持 MPS，并提示 torch.profiler 没有 MPS 的 ProfilerActivity，
算子明细只反映 CPU 侧下发耗时
- train.py / eval.py 新增 --device，配置默认改为 auto
- check_env 改为列出全部可用后端与 auto 的选择结果
实测：tiny_conv 训练在 MPS 上比 CPU 快约 32 倍。
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
203 lines
7.5 KiB
Python
203 lines
7.5 KiB
Python
"""训练循环。支持 cuda / mps / cpu 三种后端。"""
|
||
|
||
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.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 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
|