refactor: 按实际的五步工作流裁掉无关代码
实际工作流是：simclr 训练 -> 聚类 -> 人工归类 -> 分类训练 -> GradCAM/SHAP 生成
热力图与 box。据此移除不参与其中的检测/定位链路，共 31 个文件、约 2600 行：
- 模型：defect_net、postprocess、necks/、heads/center_head
- 数据：film_defect（YOLO 读取端）、target_encoder
- 引擎：trainer、evaluator、losses
- 入口：train.py、eval.py
- 脚本：convert_annotations、split_dataset、export_onnx、profile_inference
- 配置：configs/{train,eval,model,dataset}
- utils/labels（多类/单类标签模式，仅检测分类头使用）
共享模块按实际用量瘦身：engine/metrics 只留 AverageMeter（检测指标随之移除）；
datasets/yolo 只留写出侧（IMAGE_SUFFIXES 与 write_data_yaml）——第 5 步产出的
YOLO 标签是给外部检测训练用的交换格式，读取端不再需要。
裁剪过程中修掉一个真 bug：generate_boxes 的 --stage 默认 6（layer3）对 tiny_conv
这类非 nn.Sequential 骨架直接抛异常。现在自动退回骨架最终输出并警告，
另加 --stage -1 显式选择最终输出。
五步工作流已端到端实跑验证（simclr 1 轮 -> 聚类 -> 导出 -> 分类 1 轮 ->
gradcam 与 shap 两种方法各生成一遍）。
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
57446c2f10
commit
86c86b105c
+3
-7
@@ -1,4 +1,4 @@
|
||||
"""训练 / 评估流程。"""
|
||||
"""训练与分析流程：自监督、聚类、分类、归因定位。"""
|
||||
|
||||
from engine.checkpoint import load_checkpoint, save_checkpoint
|
||||
from engine.classification import (
|
||||
@@ -6,18 +6,14 @@ from engine.classification import (
|
||||
ClassificationTrainer,
|
||||
class_weights,
|
||||
)
|
||||
from engine.evaluator import Evaluator
|
||||
from engine.losses import DetectionLoss
|
||||
from engine.optim import build_optimizer, build_scheduler
|
||||
from engine.trainer import Trainer
|
||||
from engine.simclr import SimCLRTrainer
|
||||
|
||||
__all__ = [
|
||||
"Trainer",
|
||||
"ClassificationTrainer",
|
||||
"ClassificationMetric",
|
||||
"class_weights",
|
||||
"Evaluator",
|
||||
"DetectionLoss",
|
||||
"SimCLRTrainer",
|
||||
"build_optimizer",
|
||||
"build_scheduler",
|
||||
"save_checkpoint",
|
||||
|
||||
@@ -1,87 +0,0 @@
|
||||
"""评估：跑一遍验证集，输出分类与定位指标。"""
|
||||
|
||||
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
|
||||
@@ -1,79 +0,0 @@
|
||||
"""损失函数。"""
|
||||
|
||||
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())
|
||||
+4
-155
@@ -1,9 +1,10 @@
|
||||
"""指标：分类与定位。"""
|
||||
"""训练过程中的通用统计。
|
||||
|
||||
分类任务的指标（混淆矩阵、macro-F1 等）在 engine/classification.py 里。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class AverageMeter:
|
||||
def __init__(self) -> None:
|
||||
@@ -17,155 +18,3 @@ class AverageMeter:
|
||||
@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()),
|
||||
}
|
||||
|
||||
@@ -1,303 +0,0 @@
|
||||
"""训练循环。支持 cuda / mps / cpu 三种后端。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import time
|
||||
import warnings
|
||||
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 load_checkpoint, 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 resume(self, path: str | Path) -> dict:
|
||||
"""从检查点续训。
|
||||
|
||||
**配置文件是唯一真相源**：检查点只提供"进度"——模型权重、优化器动量缓冲、
|
||||
已训到第几个 epoch、global_step、best_metric；而 lr、weight_decay、总轮数、
|
||||
调度器一律以当前 yaml 为准。否则改了配置再续训，改动会被检查点静默覆盖。
|
||||
"""
|
||||
state = load_checkpoint(path, self.model, map_location=self.device)
|
||||
|
||||
if state.get("optimizer"):
|
||||
# optimizer.load_state_dict 会整体覆盖 param_groups，把 lr/weight_decay
|
||||
# 一并还原成存档时的值。先抄下配置构造出的超参，加载完按组写回，
|
||||
# 只留下动量缓冲。按组写回也保住了 BN 与 bias 那组 weight_decay=0 的结构。
|
||||
intended = [
|
||||
{key: value for key, value in group.items() if key != "params"}
|
||||
for group in self.optimizer.param_groups
|
||||
]
|
||||
self.optimizer.load_state_dict(state["optimizer"])
|
||||
for group, wanted in zip(self.optimizer.param_groups, intended, strict=True):
|
||||
group.update(wanted)
|
||||
|
||||
self.start_epoch = int(state.get("epoch", -1)) + 1
|
||||
self.global_step = int(state.get("global_step", 0))
|
||||
best = state.get("best_metric")
|
||||
self.best_metric = float(best) if best is not None and math.isfinite(float(best)) else float("-inf")
|
||||
|
||||
self._restore_scheduler()
|
||||
self._log_config_changes(state)
|
||||
logger.info(
|
||||
"已从 %s 恢复：起始 epoch=%d，best_%s=%.4f",
|
||||
path,
|
||||
self.start_epoch,
|
||||
self.monitor,
|
||||
self.best_metric,
|
||||
)
|
||||
return state
|
||||
|
||||
def _restore_scheduler(self) -> None:
|
||||
"""按**当前配置**重建学习率调度，再快进到已训到的 epoch。
|
||||
|
||||
不恢复检查点里的调度器状态：它带着当时的 T_max 与 base_lrs，一旦 yaml 里改了
|
||||
总轮数或学习率就会被静默覆盖。cosine / step / multistep 都是 epoch 的确定性函数，
|
||||
快进与恢复状态等价，所以统一走重建这一条路。
|
||||
"""
|
||||
if self.scheduler is None:
|
||||
return
|
||||
# cosine 的链式实现基于当前 lr 递推，必须先复位回 initial_lr，
|
||||
# 否则从检查点带来的 lr 出发会算出一条错的曲线。
|
||||
for group in self.optimizer.param_groups:
|
||||
group["lr"] = group.get("initial_lr", group["lr"])
|
||||
|
||||
if self.cfg.get("restart_schedule"):
|
||||
# 续训到调度末段时 cosine 会把新设的 lr 乘上一个接近 0 的系数，
|
||||
# "调大学习率接着练"就落空了。打开这个开关按剩余轮数重开一条曲线。
|
||||
remaining = max(1, self.epochs - self.start_epoch)
|
||||
self.scheduler = build_scheduler(self.optimizer, self.cfg.get("scheduler"), remaining)
|
||||
logger.warning(
|
||||
"restart_schedule=true：学习率调度从 epoch %d 重新开始，覆盖剩余 %d 轮，起始 lr=%.2e",
|
||||
self.start_epoch,
|
||||
remaining,
|
||||
self.optimizer.param_groups[0]["lr"],
|
||||
)
|
||||
return
|
||||
self.scheduler = build_scheduler(self.optimizer, self.cfg.get("scheduler"), self.epochs)
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore") # 快进阶段的 "step before optimizer.step()" 属预期
|
||||
for _ in range(self.start_epoch):
|
||||
self.scheduler.step()
|
||||
|
||||
def _log_config_changes(self, state: dict) -> None:
|
||||
"""检查点里的配置与当前配置不一致时说清楚用的是哪一份。"""
|
||||
archived = state.get("train_cfg") or {}
|
||||
changes = []
|
||||
if archived.get("epochs") not in (None, self.epochs):
|
||||
changes.append(f"总轮数 {archived['epochs']} -> {self.epochs}")
|
||||
for key, value in (archived.get("optimizer") or {}).items():
|
||||
current = dict(self.cfg.get("optimizer", {})).get(key)
|
||||
if current is not None and current != value:
|
||||
changes.append(f"optimizer.{key} {value} -> {current}")
|
||||
archived_scheduler = (archived.get("scheduler") or {}).get("name")
|
||||
current_scheduler = (dict(self.cfg.get("scheduler") or {})).get("name")
|
||||
if archived_scheduler and archived_scheduler != current_scheduler:
|
||||
changes.append(f"scheduler {archived_scheduler} -> {current_scheduler}")
|
||||
|
||||
if changes:
|
||||
logger.warning("配置与检查点不一致，按**当前配置**执行：%s", "；".join(changes))
|
||||
logger.info(
|
||||
"续训超参：lr=%.2e（initial_lr=%.2e）weight_decay=%s，共 %d 轮",
|
||||
self.optimizer.param_groups[0]["lr"],
|
||||
self.optimizer.param_groups[0].get("initial_lr", float("nan")),
|
||||
[g["weight_decay"] for g in self.optimizer.param_groups],
|
||||
self.epochs,
|
||||
)
|
||||
|
||||
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,
|
||||
scheduler=self.scheduler,
|
||||
global_step=self.global_step,
|
||||
train_cfg=dict(self.cfg),
|
||||
)
|
||||
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,
|
||||
scheduler=self.scheduler,
|
||||
global_step=self.global_step,
|
||||
train_cfg=dict(self.cfg),
|
||||
)
|
||||
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