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,17 @@
|
||||
"""训练 / 评估流程。"""
|
||||
|
||||
from engine.checkpoint import load_checkpoint, save_checkpoint
|
||||
from engine.evaluator import Evaluator
|
||||
from engine.losses import DetectionLoss
|
||||
from engine.optim import build_optimizer, build_scheduler
|
||||
from engine.trainer import Trainer
|
||||
|
||||
__all__ = [
|
||||
"Trainer",
|
||||
"Evaluator",
|
||||
"DetectionLoss",
|
||||
"build_optimizer",
|
||||
"build_scheduler",
|
||||
"save_checkpoint",
|
||||
"load_checkpoint",
|
||||
]
|
||||
@@ -0,0 +1,46 @@
|
||||
"""检查点存取。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
def save_checkpoint(
|
||||
path: str | Path,
|
||||
model: nn.Module,
|
||||
optimizer: torch.optim.Optimizer | None = None,
|
||||
epoch: int = 0,
|
||||
best_metric: float | None = None,
|
||||
config: dict | None = None,
|
||||
) -> Path:
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
torch.save(
|
||||
{
|
||||
"model": model.state_dict(),
|
||||
"optimizer": optimizer.state_dict() if optimizer else None,
|
||||
"epoch": epoch,
|
||||
"best_metric": best_metric,
|
||||
"config": config,
|
||||
},
|
||||
path,
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
def load_checkpoint(
|
||||
path: str | Path,
|
||||
model: nn.Module,
|
||||
optimizer: torch.optim.Optimizer | None = None,
|
||||
map_location: Any = "cpu",
|
||||
strict: bool = True,
|
||||
) -> dict:
|
||||
state = torch.load(Path(path), map_location=map_location, weights_only=False)
|
||||
model.load_state_dict(state["model"], strict=strict)
|
||||
if optimizer is not None and state.get("optimizer"):
|
||||
optimizer.load_state_dict(state["optimizer"])
|
||||
return state
|
||||
@@ -0,0 +1,87 @@
|
||||
"""评估：跑一遍验证集，输出分类与定位指标。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from engine.metrics import AverageMeter, DetectionMetric, MultiLabelMetric, SingleLabelMetric
|
||||
from models.postprocess import decode_heatmap
|
||||
from utils.labels import MULTILABEL
|
||||
from utils.logging import get_logger
|
||||
|
||||
logger = get_logger("engine.evaluator")
|
||||
|
||||
|
||||
class Evaluator:
|
||||
def __init__(
|
||||
self,
|
||||
model: nn.Module,
|
||||
dataloader: DataLoader,
|
||||
device: torch.device,
|
||||
criterion: nn.Module | None = None,
|
||||
score_threshold: float = 0.3,
|
||||
iou_threshold: float = 0.5,
|
||||
topk: int = 50,
|
||||
) -> None:
|
||||
self.model = model
|
||||
self.dataloader = dataloader
|
||||
self.device = device
|
||||
self.criterion = criterion
|
||||
self.score_threshold = score_threshold
|
||||
self.topk = topk
|
||||
self.det_metric = DetectionMetric(iou_threshold)
|
||||
# 分类指标按模型的 label_mode 选择：多类用多标签指标，单类用互斥分类指标。
|
||||
self.label_mode = getattr(model, "label_mode", MULTILABEL)
|
||||
if self.label_mode == MULTILABEL:
|
||||
self.cls_metric: MultiLabelMetric | SingleLabelMetric = MultiLabelMetric()
|
||||
else:
|
||||
num_outputs = model.cls_head.num_outputs if model.cls_head is not None else 1
|
||||
self.cls_metric = SingleLabelMetric(num_outputs)
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate(self) -> dict[str, float]:
|
||||
self.model.eval()
|
||||
self.det_metric.reset()
|
||||
self.cls_metric.reset()
|
||||
loss_meter = AverageMeter()
|
||||
|
||||
stride = getattr(self.model, "out_stride", 4)
|
||||
for batch in self.dataloader:
|
||||
images = batch["images"].to(self.device, non_blocking=True)
|
||||
outputs = self.model(images)
|
||||
|
||||
if self.criterion is not None:
|
||||
batch_on_device = _to_device(batch, self.device)
|
||||
losses = self.criterion(outputs, batch_on_device)
|
||||
loss_meter.update(float(losses["total"]), images.shape[0])
|
||||
|
||||
detections = decode_heatmap(
|
||||
outputs["heatmap"],
|
||||
outputs["wh"],
|
||||
outputs["offset"],
|
||||
stride=stride,
|
||||
topk=self.topk,
|
||||
score_threshold=self.score_threshold,
|
||||
)
|
||||
for i, det in enumerate(detections):
|
||||
self.det_metric.update(det["boxes"], det["labels"], batch["boxes"][i], batch["labels"][i])
|
||||
|
||||
if "logits" in outputs:
|
||||
self.cls_metric.update(outputs["logits"], batch["image_labels"])
|
||||
|
||||
metrics: dict[str, float] = {**self.det_metric.compute(), **self.cls_metric.compute()}
|
||||
if self.criterion is not None:
|
||||
metrics["loss"] = loss_meter.avg
|
||||
return metrics
|
||||
|
||||
|
||||
def _to_device(batch: dict, device: torch.device) -> dict:
|
||||
moved: dict[str, Any] = dict(batch)
|
||||
moved["images"] = batch["images"].to(device, non_blocking=True)
|
||||
moved["target"] = {k: v.to(device, non_blocking=True) for k, v in batch["target"].items()}
|
||||
moved["image_labels"] = batch["image_labels"].to(device, non_blocking=True)
|
||||
return moved
|
||||
@@ -0,0 +1,79 @@
|
||||
"""损失函数。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from utils.labels import MULTILABEL, check_label_mode
|
||||
|
||||
|
||||
def gaussian_focal_loss(
|
||||
pred_logits: torch.Tensor,
|
||||
gt_heatmap: torch.Tensor,
|
||||
alpha: float = 2.0,
|
||||
beta: float = 4.0,
|
||||
eps: float = 1e-4,
|
||||
) -> torch.Tensor:
|
||||
"""CenterNet 的 penalty-reduced focal loss，按正样本数归一化。"""
|
||||
pred = pred_logits.sigmoid().clamp(eps, 1 - eps)
|
||||
pos_mask = gt_heatmap.eq(1.0).float()
|
||||
neg_mask = 1.0 - pos_mask
|
||||
|
||||
pos_loss = -((1 - pred) ** alpha) * torch.log(pred) * pos_mask
|
||||
neg_loss = -((1 - gt_heatmap) ** beta) * (pred**alpha) * torch.log(1 - pred) * neg_mask
|
||||
|
||||
num_pos = pos_mask.sum().clamp(min=1.0)
|
||||
return (pos_loss.sum() + neg_loss.sum()) / num_pos
|
||||
|
||||
|
||||
def masked_l1_loss(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
|
||||
"""只在中心点位置计算 L1。mask: (B, H, W)。"""
|
||||
mask = mask.unsqueeze(1).expand_as(pred)
|
||||
loss = F.l1_loss(pred * mask, target * mask, reduction="sum")
|
||||
return loss / mask.sum().clamp(min=1.0)
|
||||
|
||||
|
||||
class DetectionLoss(nn.Module):
|
||||
"""heatmap + wh + offset (+ 图像级分类) 的加权和。
|
||||
|
||||
分类分支按 `label_mode` 切换：
|
||||
multilabel（多类）—— 多热标签，逐类 BCE；
|
||||
multiclass（单类）—— 互斥下标，softmax 交叉熵。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
heatmap_weight: float = 1.0,
|
||||
wh_weight: float = 0.1,
|
||||
offset_weight: float = 1.0,
|
||||
cls_weight: float = 1.0,
|
||||
label_mode: str = MULTILABEL,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.label_mode = check_label_mode(label_mode)
|
||||
self.heatmap_weight = heatmap_weight
|
||||
self.wh_weight = wh_weight
|
||||
self.offset_weight = offset_weight
|
||||
self.cls_weight = cls_weight
|
||||
|
||||
def forward(self, outputs: dict, batch: dict) -> dict[str, torch.Tensor]:
|
||||
target = batch["target"]
|
||||
losses = {
|
||||
"heatmap": self.heatmap_weight * gaussian_focal_loss(outputs["heatmap"], target["heatmap"]),
|
||||
"wh": self.wh_weight * masked_l1_loss(outputs["wh"], target["wh"], target["mask"]),
|
||||
"offset": self.offset_weight
|
||||
* masked_l1_loss(outputs["offset"], target["offset"], target["mask"]),
|
||||
}
|
||||
if "logits" in outputs and self.cls_weight > 0:
|
||||
losses["cls"] = self.cls_weight * self._classification_loss(
|
||||
outputs["logits"], batch["image_labels"]
|
||||
)
|
||||
losses["total"] = sum(losses.values())
|
||||
return losses
|
||||
|
||||
def _classification_loss(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||||
if self.label_mode == MULTILABEL:
|
||||
return F.binary_cross_entropy_with_logits(logits, targets.float())
|
||||
return F.cross_entropy(logits, targets.long())
|
||||
@@ -0,0 +1,171 @@
|
||||
"""指标：分类与定位。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class AverageMeter:
|
||||
def __init__(self) -> None:
|
||||
self.total = 0.0
|
||||
self.count = 0
|
||||
|
||||
def update(self, value: float, n: int = 1) -> None:
|
||||
self.total += float(value) * n
|
||||
self.count += n
|
||||
|
||||
@property
|
||||
def avg(self) -> float:
|
||||
return self.total / self.count if self.count else 0.0
|
||||
|
||||
|
||||
def box_iou(boxes_a: torch.Tensor, boxes_b: torch.Tensor) -> torch.Tensor:
|
||||
"""(N, 4) 与 (M, 4) 的两两 IoU，返回 (N, M)。"""
|
||||
if boxes_a.numel() == 0 or boxes_b.numel() == 0:
|
||||
return boxes_a.new_zeros((boxes_a.shape[0], boxes_b.shape[0]))
|
||||
area_a = (boxes_a[:, 2] - boxes_a[:, 0]).clamp(min=0) * (boxes_a[:, 3] - boxes_a[:, 1]).clamp(min=0)
|
||||
area_b = (boxes_b[:, 2] - boxes_b[:, 0]).clamp(min=0) * (boxes_b[:, 3] - boxes_b[:, 1]).clamp(min=0)
|
||||
|
||||
lt = torch.max(boxes_a[:, None, :2], boxes_b[None, :, :2])
|
||||
rb = torch.min(boxes_a[:, None, 2:], boxes_b[None, :, 2:])
|
||||
wh = (rb - lt).clamp(min=0)
|
||||
inter = wh[..., 0] * wh[..., 1]
|
||||
return inter / (area_a[:, None] + area_b[None, :] - inter).clamp(min=1e-6)
|
||||
|
||||
|
||||
class DetectionMetric:
|
||||
"""按 IoU 阈值做贪心匹配，统计 precision / recall / F1 与平均 IoU。
|
||||
|
||||
这是工程上够用的近似指标；正式对比模型时建议再跑一次 COCO mAP。
|
||||
"""
|
||||
|
||||
def __init__(self, iou_threshold: float = 0.5) -> None:
|
||||
self.iou_threshold = iou_threshold
|
||||
self.reset()
|
||||
|
||||
def reset(self) -> None:
|
||||
self.tp = 0
|
||||
self.fp = 0
|
||||
self.fn = 0
|
||||
self.iou_sum = 0.0
|
||||
|
||||
def update(
|
||||
self,
|
||||
pred_boxes: torch.Tensor,
|
||||
pred_labels: torch.Tensor,
|
||||
gt_boxes: torch.Tensor,
|
||||
gt_labels: torch.Tensor,
|
||||
) -> None:
|
||||
if gt_boxes.numel() == 0:
|
||||
self.fp += int(pred_boxes.shape[0])
|
||||
return
|
||||
if pred_boxes.numel() == 0:
|
||||
self.fn += int(gt_boxes.shape[0])
|
||||
return
|
||||
|
||||
ious = box_iou(pred_boxes.float().cpu(), gt_boxes.float().cpu())
|
||||
same_class = pred_labels.cpu()[:, None] == gt_labels.cpu()[None, :]
|
||||
ious = ious * same_class.float()
|
||||
|
||||
matched_gt: set[int] = set()
|
||||
for pred_idx in range(ious.shape[0]):
|
||||
best_iou, best_gt = ious[pred_idx].max(dim=0)
|
||||
gt_idx = int(best_gt)
|
||||
if float(best_iou) >= self.iou_threshold and gt_idx not in matched_gt:
|
||||
matched_gt.add(gt_idx)
|
||||
self.tp += 1
|
||||
self.iou_sum += float(best_iou)
|
||||
else:
|
||||
self.fp += 1
|
||||
self.fn += int(gt_boxes.shape[0]) - len(matched_gt)
|
||||
|
||||
def compute(self) -> dict[str, float]:
|
||||
precision = self.tp / max(self.tp + self.fp, 1)
|
||||
recall = self.tp / max(self.tp + self.fn, 1)
|
||||
f1 = 2 * precision * recall / max(precision + recall, 1e-9)
|
||||
return {
|
||||
"precision": precision,
|
||||
"recall": recall,
|
||||
"f1": f1,
|
||||
"mean_iou": self.iou_sum / max(self.tp, 1),
|
||||
}
|
||||
|
||||
|
||||
class MultiLabelMetric:
|
||||
"""多类模式：逐类阈值化后统计 micro precision / recall / F1 与整图完全正确率。"""
|
||||
|
||||
def __init__(self, threshold: float = 0.5) -> None:
|
||||
self.threshold = threshold
|
||||
self.reset()
|
||||
|
||||
def reset(self) -> None:
|
||||
self.tp = 0
|
||||
self.fp = 0
|
||||
self.fn = 0
|
||||
self.exact = 0
|
||||
self.total = 0
|
||||
|
||||
def update(self, logits: torch.Tensor, targets: torch.Tensor) -> None:
|
||||
pred = (logits.sigmoid() >= self.threshold).float().cpu()
|
||||
gt = targets.float().cpu()
|
||||
self.tp += int((pred * gt).sum())
|
||||
self.fp += int((pred * (1 - gt)).sum())
|
||||
self.fn += int(((1 - pred) * gt).sum())
|
||||
self.exact += int((pred == gt).all(dim=1).sum())
|
||||
self.total += int(gt.shape[0])
|
||||
|
||||
def compute(self) -> dict[str, float]:
|
||||
precision = self.tp / max(self.tp + self.fp, 1)
|
||||
recall = self.tp / max(self.tp + self.fn, 1)
|
||||
return {
|
||||
"cls_precision": precision,
|
||||
"cls_recall": recall,
|
||||
"cls_f1": 2 * precision * recall / max(precision + recall, 1e-9),
|
||||
"cls_exact_match": self.exact / max(self.total, 1),
|
||||
}
|
||||
|
||||
|
||||
class SingleLabelMetric:
|
||||
"""单类模式：互斥分类的准确率与 macro precision / recall / F1。
|
||||
|
||||
指标名与 `MultiLabelMetric` 保持一致（cls_precision / cls_recall / cls_f1），
|
||||
因此切换 label_mode 后 `train.monitor` 不需要改。
|
||||
"""
|
||||
|
||||
def __init__(self, num_outputs: int) -> None:
|
||||
self.num_outputs = num_outputs
|
||||
self.reset()
|
||||
|
||||
def reset(self) -> None:
|
||||
self.tp = torch.zeros(self.num_outputs)
|
||||
self.fp = torch.zeros(self.num_outputs)
|
||||
self.fn = torch.zeros(self.num_outputs)
|
||||
self.correct = 0
|
||||
self.total = 0
|
||||
|
||||
def update(self, logits: torch.Tensor, targets: torch.Tensor) -> None:
|
||||
pred = logits.argmax(dim=1).cpu()
|
||||
gt = targets.long().cpu()
|
||||
for cls in range(self.num_outputs):
|
||||
pred_cls = pred == cls
|
||||
gt_cls = gt == cls
|
||||
self.tp[cls] += int((pred_cls & gt_cls).sum())
|
||||
self.fp[cls] += int((pred_cls & ~gt_cls).sum())
|
||||
self.fn[cls] += int((~pred_cls & gt_cls).sum())
|
||||
self.correct += int((pred == gt).sum())
|
||||
self.total += int(gt.numel())
|
||||
|
||||
def compute(self) -> dict[str, float]:
|
||||
# 只对验证集中出现过的类别做 macro 平均，避免缺席类别把指标拉到 0。
|
||||
present = (self.tp + self.fn) > 0
|
||||
precision = self.tp / (self.tp + self.fp).clamp(min=1)
|
||||
recall = self.tp / (self.tp + self.fn).clamp(min=1)
|
||||
f1 = 2 * precision * recall / (precision + recall).clamp(min=1e-9)
|
||||
if present.any():
|
||||
precision, recall, f1 = precision[present], recall[present], f1[present]
|
||||
return {
|
||||
"cls_accuracy": self.correct / max(self.total, 1),
|
||||
"cls_precision": float(precision.mean()),
|
||||
"cls_recall": float(recall.mean()),
|
||||
"cls_f1": float(f1.mean()),
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
"""优化器与学习率调度。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
def build_optimizer(model: nn.Module, cfg: Mapping[str, Any]) -> torch.optim.Optimizer:
|
||||
"""`cfg` 形如 `{name: sgd, lr: 0.01, momentum: 0.9, weight_decay: 1e-4}`。
|
||||
|
||||
BatchNorm 与 bias 不做 weight decay——小数据集上这项影响明显。
|
||||
"""
|
||||
cfg = dict(cfg)
|
||||
name = str(cfg.pop("name", "adamw")).lower()
|
||||
weight_decay = float(cfg.pop("weight_decay", 0.0))
|
||||
|
||||
decay, no_decay = [], []
|
||||
for module in model.modules():
|
||||
for param_name, param in module.named_parameters(recurse=False):
|
||||
if not param.requires_grad:
|
||||
continue
|
||||
if param_name == "bias" or isinstance(module, nn.modules.batchnorm._BatchNorm):
|
||||
no_decay.append(param)
|
||||
else:
|
||||
decay.append(param)
|
||||
groups = [
|
||||
{"params": decay, "weight_decay": weight_decay},
|
||||
{"params": no_decay, "weight_decay": 0.0},
|
||||
]
|
||||
|
||||
if name == "sgd":
|
||||
cfg.setdefault("momentum", 0.9)
|
||||
cfg.setdefault("nesterov", True)
|
||||
return torch.optim.SGD(groups, **cfg)
|
||||
if name == "adam":
|
||||
return torch.optim.Adam(groups, **cfg)
|
||||
if name == "adamw":
|
||||
return torch.optim.AdamW(groups, **cfg)
|
||||
raise ValueError(f"未知优化器: {name}")
|
||||
|
||||
|
||||
def build_scheduler(optimizer: torch.optim.Optimizer, cfg: Mapping[str, Any] | None, epochs: int):
|
||||
"""返回按 epoch 步进的 scheduler；`cfg` 为空则不使用。"""
|
||||
if not cfg:
|
||||
return None
|
||||
cfg = dict(cfg)
|
||||
name = str(cfg.pop("name", "cosine")).lower()
|
||||
if name == "cosine":
|
||||
cfg.setdefault("T_max", epochs)
|
||||
return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, **cfg)
|
||||
if name == "step":
|
||||
cfg.setdefault("step_size", max(epochs // 3, 1))
|
||||
return torch.optim.lr_scheduler.StepLR(optimizer, **cfg)
|
||||
if name == "multistep":
|
||||
cfg.setdefault("milestones", [int(epochs * 0.6), int(epochs * 0.85)])
|
||||
return torch.optim.lr_scheduler.MultiStepLR(optimizer, **cfg)
|
||||
raise ValueError(f"未知调度器: {name}")
|
||||
@@ -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
|
||||
@@ -0,0 +1,100 @@
|
||||
"""训练可视化：SwanLab 封装。
|
||||
|
||||
配置示例（`train.swanlab`）：
|
||||
|
||||
```yaml
|
||||
train:
|
||||
swanlab:
|
||||
enabled: true
|
||||
project: film-defect
|
||||
name: baseline # 缺省用 work_dir 的目录名
|
||||
mode: online # online(云端) | local(本地看板) | offline(仅存盘) | disabled
|
||||
tags: [tiny_conv, center-head]
|
||||
log_dir: null # 缺省 <work_dir>/swanlog
|
||||
```
|
||||
|
||||
未安装 swanlab、或 `enabled: false` 时整体退化为空操作，训练流程不受影响。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from utils.logging import get_logger
|
||||
|
||||
logger = get_logger("engine.visualizer")
|
||||
|
||||
VALID_MODES = ("online", "local", "offline", "disabled")
|
||||
|
||||
|
||||
class SwanLabLogger:
|
||||
"""把标量指标写到 SwanLab；不可用时静默降级。
|
||||
|
||||
Args:
|
||||
enabled: 关掉后所有方法都是空操作。
|
||||
work_dir: 用于推导默认的实验名与 `log_dir`。
|
||||
config: 记录到实验里的超参（一般直接传完整 yaml 配置）。
|
||||
其余参数透传给 `swanlab.init`。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
enabled: bool = True,
|
||||
work_dir: str | Path = "runs",
|
||||
config: dict | None = None,
|
||||
project: str = "film-defect",
|
||||
name: str | None = None,
|
||||
mode: str = "online",
|
||||
log_dir: str | Path | None = None,
|
||||
**init_kwargs: Any,
|
||||
) -> None:
|
||||
work_dir = Path(work_dir)
|
||||
self.run = None
|
||||
self.enabled = bool(enabled) and mode != "disabled"
|
||||
if not self.enabled:
|
||||
logger.info("SwanLab 可视化未启用")
|
||||
return
|
||||
|
||||
if mode not in VALID_MODES:
|
||||
raise ValueError(f"未知 swanlab mode: {mode}，可选 {VALID_MODES}")
|
||||
|
||||
try:
|
||||
import swanlab
|
||||
except ImportError:
|
||||
logger.warning("未安装 swanlab，跳过可视化（uv sync 后可用；或设 train.swanlab.enabled=false）")
|
||||
self.enabled = False
|
||||
return
|
||||
|
||||
try:
|
||||
self.run = swanlab.init(
|
||||
project=project,
|
||||
name=name or work_dir.name,
|
||||
mode=mode,
|
||||
config=config or {},
|
||||
log_dir=str(log_dir or work_dir / "swanlog"),
|
||||
**init_kwargs,
|
||||
)
|
||||
self._swanlab = swanlab
|
||||
logger.info("SwanLab 已启动（project=%s, mode=%s）", project, mode)
|
||||
except Exception as exc: # 联网失败 / 未登录不应中断训练
|
||||
logger.warning("SwanLab 初始化失败，训练继续但不记录可视化: %s", exc)
|
||||
self.enabled = False
|
||||
|
||||
def log(self, metrics: dict[str, float], step: int | None = None, prefix: str = "") -> None:
|
||||
if not self.enabled:
|
||||
return
|
||||
payload = {f"{prefix}{key}": float(value) for key, value in metrics.items()}
|
||||
try:
|
||||
self._swanlab.log(payload, step=step)
|
||||
except Exception as exc: # 记录失败同样不中断训练
|
||||
logger.warning("SwanLab 记录失败: %s", exc)
|
||||
self.enabled = False
|
||||
|
||||
def finish(self) -> None:
|
||||
if not self.enabled:
|
||||
return
|
||||
try:
|
||||
self._swanlab.finish()
|
||||
except Exception as exc:
|
||||
logger.warning("SwanLab 收尾失败: %s", exc)
|
||||
Reference in New Issue
Block a user