Shaun ZhouandClaude Opus 5
33fdf396ee
在聚类 + 人工归类得到的数据集（data/<集>/<类别名>/*.png）上训练分类模型。 - datasets/classification.py：按目录组织的分类数据集，**逐类分层**切分 train/val （others 只有 65 张，整体随机切很容易让它在验证集里挂零），固定种子可复现 - models/classifier.py：卷积骨架 + 分类头，与 DefectNet 共用 backbone 注册表， 因此自监督产出的 simclr_backbone.pth 可用 init_from 直接初始化； frozen_backbone 可做线性探针 - engine/classification.py：训练循环 + 指标。主指标用 macro-F1 而非准确率—— 数据从 normal 2679 张到 others 65 张差 41 倍，全预测成 normal 也有 24% 准确率； 损失按类别频率加权，避免小类被淹没；输出逐类 P/R/F1 与混淆矩阵 - classify.py：训练/评估入口，--eval-only 可单独评估并导出 eval_report.json - configs/classify/film.yaml：预处理沿用聚类阶段的结论（拉伸 + 逐图标准化）。 抹掉亮度在这里尤其重要：有几类几乎只来自单一原始数据集，不抹掉的话模型可能 靠"认出这是哪个数据集"作弊 - 续训语义与检测/自监督一致：检查点只提供进度，超参以配置为准 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
289 lines
12 KiB
Python
289 lines
12 KiB
Python
"""分类任务的训练、评估与指标。
|
||
|
||
数据极不均衡（normal 2679 张 vs others 65 张，差 41 倍），所以：
|
||
- 主指标用 **macro-F1** 而不是准确率——全预测成 normal 也能拿到 24% 准确率，但 macro-F1 会很低；
|
||
- 损失支持按类别频率加权，避免小类被淹没。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
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.nn import functional as F
|
||
from torch.utils.data import DataLoader
|
||
|
||
from engine.checkpoint import load_checkpoint, save_checkpoint
|
||
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.logging import get_logger
|
||
|
||
logger = get_logger("engine.classification")
|
||
|
||
|
||
class ClassificationMetric:
|
||
"""混淆矩阵 + 由它导出的各项指标。
|
||
|
||
准确率在不均衡数据上会骗人，所以同时给 macro-F1（各类等权）与
|
||
weighted-F1（按样本数加权），并保留逐类明细供人工判断哪些类没学会。
|
||
"""
|
||
|
||
def __init__(self, classes: list[str]) -> None:
|
||
self.classes = list(classes)
|
||
self.reset()
|
||
|
||
def reset(self) -> None:
|
||
n = len(self.classes)
|
||
self.confusion = torch.zeros(n, n, dtype=torch.long) # 行=真实，列=预测
|
||
|
||
def update(self, logits: torch.Tensor, targets: torch.Tensor) -> None:
|
||
preds = logits.argmax(dim=1).cpu()
|
||
for target, pred in zip(targets.cpu().tolist(), preds.tolist(), strict=True):
|
||
self.confusion[target, pred] += 1
|
||
|
||
def per_class(self) -> dict[str, dict[str, float]]:
|
||
report = {}
|
||
for index, name in enumerate(self.classes):
|
||
tp = int(self.confusion[index, index])
|
||
support = int(self.confusion[index].sum())
|
||
predicted = int(self.confusion[:, index].sum())
|
||
precision = tp / predicted if predicted else 0.0
|
||
recall = tp / support if support else 0.0
|
||
f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0
|
||
report[name] = {
|
||
"precision": precision,
|
||
"recall": recall,
|
||
"f1": f1,
|
||
"support": support,
|
||
}
|
||
return report
|
||
|
||
def compute(self) -> dict[str, float]:
|
||
total = int(self.confusion.sum())
|
||
if not total:
|
||
return {"accuracy": 0.0, "macro_f1": 0.0, "weighted_f1": 0.0, "balanced_accuracy": 0.0}
|
||
report = self.per_class()
|
||
f1s = [r["f1"] for r in report.values()]
|
||
recalls = [r["recall"] for r in report.values() if r["support"]]
|
||
supports = [r["support"] for r in report.values()]
|
||
weighted = sum(r["f1"] * r["support"] for r in report.values()) / total
|
||
return {
|
||
"accuracy": float(self.confusion.diagonal().sum()) / total,
|
||
"macro_f1": sum(f1s) / len(f1s),
|
||
"weighted_f1": weighted,
|
||
# 各类召回的均值：小类分错会明显拉低它，比准确率更能反映"每类都学会了吗"
|
||
"balanced_accuracy": sum(recalls) / len(recalls) if recalls else 0.0,
|
||
"num_classes_with_support": sum(1 for s in supports if s),
|
||
}
|
||
|
||
|
||
def class_weights(counts: list[int], mode: str | None = "balanced") -> torch.Tensor | None:
|
||
"""按类别频率算损失权重。
|
||
|
||
`balanced` 用逆频率并归一化到均值 1：小类的每个样本被放大，
|
||
否则 normal（占 24%）会主导梯度，小类学不动。
|
||
"""
|
||
if not mode or str(mode).lower() in ("none", "false"):
|
||
return None
|
||
if str(mode).lower() != "balanced":
|
||
raise ValueError(f"未知 class_weight: {mode}，可选 balanced / none")
|
||
counts_tensor = torch.tensor(counts, dtype=torch.float32).clamp(min=1)
|
||
weights = counts_tensor.sum() / (len(counts) * counts_tensor)
|
||
return weights / weights.mean()
|
||
|
||
|
||
class ClassificationTrainer:
|
||
"""Args:
|
||
model: 已 `.to(device)` 的分类模型。
|
||
train_loader / val_loader: 数据加载器。
|
||
cfg: `train` 段配置（epochs / optimizer / scheduler / amp / class_weight / monitor 等）。
|
||
classes: 类别名，顺序必须与标签下标一致。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
model: nn.Module,
|
||
train_loader: DataLoader,
|
||
val_loader: DataLoader | None,
|
||
cfg: Mapping[str, Any],
|
||
device: torch.device,
|
||
work_dir: str | Path,
|
||
classes: list[str],
|
||
train_counts: list[int] | None = None,
|
||
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.classes = list(classes)
|
||
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", 30))
|
||
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", "macro_f1"))
|
||
self.label_smoothing = float(self.cfg.get("label_smoothing", 0.0))
|
||
|
||
weights = class_weights(train_counts or [], self.cfg.get("class_weight", "balanced"))
|
||
self.class_weight = weights.to(device) if weights is not None else None
|
||
if self.class_weight is not None:
|
||
extremes = sorted(zip(self.classes, weights.tolist(), strict=True), key=lambda kv: -kv[1])
|
||
logger.info(
|
||
"类别权重（逆频率）：最大 %s=%.2f，最小 %s=%.2f",
|
||
extremes[0][0], extremes[0][1], extremes[-1][0], extremes[-1][1],
|
||
)
|
||
|
||
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 = resolve_amp(self.cfg.get("amp", "auto"), device)
|
||
self.scaler = torch.amp.GradScaler(device.type, enabled=self.amp)
|
||
self.peak_memory = PeakMemory(device)
|
||
|
||
self.metric = ClassificationMetric(self.classes)
|
||
self.best_metric = float("-inf")
|
||
self.start_epoch = 0
|
||
self.global_step = 0
|
||
self.visualizer = SwanLabLogger(
|
||
work_dir=self.work_dir, config=full_config, **dict(self.cfg.get("swanlab", {}))
|
||
)
|
||
|
||
def _loss(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||
return F.cross_entropy(
|
||
logits, targets, weight=self.class_weight, label_smoothing=self.label_smoothing
|
||
)
|
||
|
||
def train_one_epoch(self, epoch: int) -> dict[str, float]:
|
||
self.model.train()
|
||
loss_meter, acc_meter = AverageMeter(), AverageMeter()
|
||
start = time.time()
|
||
|
||
for step, batch in enumerate(self.train_loader):
|
||
images = batch["images"].to(self.device, non_blocking=True)
|
||
targets = batch["labels"].to(self.device, non_blocking=True)
|
||
|
||
with torch.amp.autocast(self.device.type, enabled=self.amp):
|
||
logits = self.model(images)
|
||
loss = self._loss(logits, targets)
|
||
|
||
self.optimizer.zero_grad(set_to_none=True)
|
||
self.scaler.scale(loss).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()
|
||
|
||
self.peak_memory.update()
|
||
batch_size = images.shape[0]
|
||
loss_meter.update(float(loss.detach()), batch_size)
|
||
acc_meter.update(float((logits.argmax(1) == targets).float().mean()), batch_size)
|
||
self.global_step += 1
|
||
|
||
if step % self.log_interval == 0:
|
||
logger.info(
|
||
"epoch %d [%d/%d] loss=%.4f acc=%.4f lr=%.2e",
|
||
epoch, step, len(self.train_loader), loss_meter.avg, acc_meter.avg,
|
||
self.optimizer.param_groups[0]["lr"],
|
||
)
|
||
self.visualizer.log(
|
||
{"loss": loss_meter.avg, "acc": acc_meter.avg,
|
||
"lr": self.optimizer.param_groups[0]["lr"]},
|
||
step=self.global_step, prefix="train/",
|
||
)
|
||
|
||
logger.info("epoch %d 训练耗时 %.1fs", epoch, time.time() - start)
|
||
return {"loss": loss_meter.avg, "acc": acc_meter.avg}
|
||
|
||
@torch.no_grad()
|
||
def evaluate(self, loader: DataLoader | None = None) -> dict[str, float]:
|
||
loader = loader or self.val_loader
|
||
if loader is None:
|
||
return {}
|
||
self.model.eval()
|
||
self.metric.reset()
|
||
loss_meter = AverageMeter()
|
||
|
||
for batch in loader:
|
||
images = batch["images"].to(self.device, non_blocking=True)
|
||
targets = batch["labels"].to(self.device, non_blocking=True)
|
||
logits = self.model(images)
|
||
loss_meter.update(float(self._loss(logits, targets)), images.shape[0])
|
||
self.metric.update(logits, targets)
|
||
|
||
metrics = self.metric.compute()
|
||
metrics["loss"] = loss_meter.avg
|
||
return metrics
|
||
|
||
def fit(self) -> float:
|
||
logger.info("开始训练：共 %d 个 epoch，%d 类，输出目录 %s",
|
||
self.epochs, len(self.classes), self.work_dir)
|
||
for epoch in range(self.start_epoch, self.epochs):
|
||
self.train_one_epoch(epoch)
|
||
if self.scheduler is not None:
|
||
self.scheduler.step()
|
||
|
||
metrics = self.evaluate()
|
||
if metrics:
|
||
logger.info(
|
||
"epoch %d 验证: %s",
|
||
epoch,
|
||
" ".join(f"{k}={v:.4f}" for k, v in metrics.items() if isinstance(v, float)),
|
||
)
|
||
self.visualizer.log(metrics, step=self.global_step, prefix="val/")
|
||
|
||
self._save("last.pth", epoch)
|
||
score = metrics.get(self.monitor)
|
||
if score is not None and score > self.best_metric:
|
||
self.best_metric = score
|
||
self._save("best.pth", epoch)
|
||
logger.info("epoch %d 刷新最佳 %s=%.4f，已保存 best.pth", epoch, self.monitor, score)
|
||
|
||
self.visualizer.finish()
|
||
return self.best_metric
|
||
|
||
def _save(self, name: str, epoch: int) -> None:
|
||
save_checkpoint(
|
||
self.work_dir / name, 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),
|
||
)
|
||
|
||
def resume(self, path: str | Path) -> dict:
|
||
"""续训。与检测训练同样的语义：检查点只提供进度，超参以配置为准。"""
|
||
state = load_checkpoint(path, self.model, map_location=self.device)
|
||
if state.get("optimizer"):
|
||
intended = [
|
||
{k: v for k, v in group.items() if k != "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 else float("-inf")
|
||
|
||
if self.scheduler is not None:
|
||
for group in self.optimizer.param_groups:
|
||
group["lr"] = group.get("initial_lr", group["lr"])
|
||
self.scheduler = build_scheduler(self.optimizer, self.cfg.get("scheduler"), self.epochs)
|
||
with warnings.catch_warnings():
|
||
warnings.simplefilter("ignore")
|
||
for _ in range(self.start_epoch):
|
||
self.scheduler.step()
|
||
logger.info("已从 %s 恢复：起始 epoch=%d，best_%s=%.4f",
|
||
path, self.start_epoch, self.monitor, self.best_metric)
|
||
return state
|