feat: 新增缺陷分类训练与评估链路

在聚类 + 人工归类得到的数据集（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>

README.md

@@ -18,7 +18,7 @@ convolution 的结构（不含 attention / transformer），并支持导出 ONNX
 | `docs/` | 算法与数学说明文档 |
 | `figures/` | 文档用的 mermaid 图（`.mmd` 源 + `.md` 预览 + `.png`） |
 
-入口在仓库根目录：[train.py](train.py)、[eval.py](eval.py)。
+入口在仓库根目录：[train.py](train.py)（检测+定位训练）、[eval.py](eval.py)、[classify.py](classify.py)（分类训练与评估）。
 
 ## 环境
 
@@ -585,6 +585,53 @@ uv run swanlab watch runs/simclr/swanlog
 训练出的骨架和检测模型用的是同一套卷积骨架，可以直接拿去初始化 `DefectNet` 的 backbone——
 标注量有限时，通常比 ImageNet 权重更贴合薄膜的成像特性。
 
+## 缺陷分类模型
+
+在聚类 + 人工归类得到的数据集上训练分类模型（`data/<数据集>/<类别名>/*.png`）：
+
+```bash
+uv run python classify.py --config configs/classify/film.yaml
+
+# 只评估
+uv run python classify.py --config configs/classify/film.yaml --eval-only \
+    --checkpoint runs/classify/best.pth
+```
+
+产物在 `work_dir`：`best.pth` / `last.pth` / `config.yaml` / `classify.log` /
+`eval_report.json`（逐类 P/R/F1 与完整混淆矩阵）。
+
+### 三个关键取舍
+
+**主指标是 macro-F1，不是准确率。** 数据从 `normal` 2679 张到 `others` 65 张差 41 倍，
+全预测成 normal 也能拿 24% 准确率。配置里同时给出 `balanced_accuracy`（各类召回的均值），
+两者一起看才知道小类有没有学会。
+
+**损失按类别频率加权**（`class_weight: balanced`，逆频率归一化到均值 1），
+否则占 24% 的 normal 会主导梯度。不需要时设 `none`。
+
+**逐图标准化抹掉整体亮度。** 15 个类别里有几类几乎只来自单一原始数据集
+（`ripple` 几乎全是 pearlescent，`grease`/`zigzag` 全是 film），不抹掉亮度的话模型
+可能靠"认出这是哪个数据集"就把类别猜对——验证集上分数很高，换批数据就崩。
+
+### 划分
+
+没有预先的 train/val，按**类别分层**切分（默认 8:2，固定种子）。逐类切分是必须的：
+`others` 只有 65 张，整体随机切很容易让它在验证集里一张不剩。
+
+### 用自监督骨架初始化
+
+前面自监督训练出的骨架可以直接拿来初始化——这正是聚类阶段产物的下游用途：
+
+```yaml
+model:
+  backbone: {name: torchvision, arch: resnet101}   # 必须与自监督时一致
+  init_from: runs/simclr_r101/simclr_backbone.pth
+  frozen_backbone: false    # true = 冻结骨架只训分类头（线性探针）
+```
+
+只加载 backbone，分类头仍随机初始化。`frozen_backbone: true` 可以单独衡量
+自监督特征本身有多好。
+
 ## 方法
 
 - **骨架**：`tiny_conv`（自研深度可分离卷积，默认）或 `torchvision`（resnet / mobilenet / shufflenet）。

classify.py

@@ -0,0 +1,180 @@
+#!/usr/bin/env python
+"""缺陷分类模型的训练与评估。
+
+    uv run python classify.py --config configs/classify/film.yaml
+    uv run python classify.py --config configs/classify/film.yaml --eval-only \
+        --checkpoint runs/classify/best.pth
+
+数据是 scripts/export_clusters.py 的产物：`root/<类别名>/*.png`。
+没有预先划分 train/val，按类别分层切分（固定种子，两次运行划分一致）。
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+from datasets import FolderClassificationDataset, classification_collate
+from engine import ClassificationTrainer
+from models import build_model
+from utils import load_config, seed_everything, setup_logging
+from utils.device import DeviceNotAvailableError, describe_device, resolve_device
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="薄膜缺陷分类模型训练/评估")
+    parser.add_argument("--config", required=True, help="yaml 配置路径")
+    parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE", help="覆盖配置项")
+    parser.add_argument("--work-dir", default=None)
+    parser.add_argument("--device", default=None, help="auto / cuda / mps / cpu")
+    parser.add_argument("--resume", default=None, help="从检查点续训")
+    parser.add_argument("--eval-only", action="store_true", help="只评估，不训练")
+    parser.add_argument("--checkpoint", default=None, help="评估用的权重")
+    parser.add_argument("--report", default=None, help="把逐类指标与混淆矩阵写成 json")
+    return parser.parse_args()
+
+
+def build_loaders(cfg, split_cfg, classes=None):
+    from torch.utils.data import DataLoader
+
+    datasets = {}
+    for split, transforms_key in (("train", "train_transforms"), ("val", "val_transforms")):
+        datasets[split] = FolderClassificationDataset(
+            root=split_cfg["root"],
+            split=split,
+            val_ratio=float(split_cfg.get("val_ratio", 0.2)),
+            seed=int(cfg.get("seed", 42)),
+            classes=classes or split_cfg.get("classes"),
+            transforms=split_cfg[transforms_key],
+            exclude=split_cfg.get("exclude") or [],
+        )
+    loader_cfg = dict(split_cfg.get("loader", {}))
+    loaders = {
+        split: DataLoader(
+            dataset,
+            batch_size=int(loader_cfg.get("batch_size", 64)),
+            shuffle=(split == "train"),
+            num_workers=int(loader_cfg.get("num_workers", 4)),
+            collate_fn=classification_collate,
+            drop_last=False,
+        )
+        for split, dataset in datasets.items()
+    }
+    return datasets, loaders
+
+
+def write_report(path: Path, trainer, metrics: dict) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(
+        json.dumps(
+            {
+                "classes": trainer.classes,
+                "metrics": metrics,
+                "per_class": trainer.metric.per_class(),
+                "confusion_matrix": trainer.metric.confusion.tolist(),
+            },
+            ensure_ascii=False,
+            indent=2,
+        ),
+        encoding="utf-8",
+    )
+
+
+def log_per_class(logger, trainer) -> None:
+    logger.info("%-14s %9s %8s %8s %8s", "类别", "precision", "recall", "f1", "support")
+    for name, row in sorted(trainer.metric.per_class().items(), key=lambda kv: -kv[1]["support"]):
+        logger.info(
+            "%-14s %9.3f %8.3f %8.3f %8d",
+            name, row["precision"], row["recall"], row["f1"], row["support"],
+        )
+
+
+def main() -> int:
+    args = parse_args()
+    cfg = load_config(args.config, args.set)
+
+    work_dir = Path(args.work_dir or cfg.get("work_dir", "runs/classify"))
+    work_dir.mkdir(parents=True, exist_ok=True)
+    logger = setup_logging(work_dir / "classify.log")
+
+    try:
+        device = resolve_device(args.device or cfg.get("device", "auto"))
+    except DeviceNotAvailableError as exc:
+        logger.error("%s", exc)
+        return 1
+    logger.info("设备: %s", describe_device(device))
+    seed_everything(int(cfg.get("seed", 42)))
+
+    datasets, loaders = build_loaders(cfg, cfg["data"])
+    classes = datasets["train"].classes
+    counts = datasets["train"].class_counts()
+    logger.info(
+        "数据集 %s：train %d / val %d，共 %d 类",
+        cfg["data"]["root"], len(datasets["train"]), len(datasets["val"]), len(classes),
+    )
+    logger.info("训练集各类样本数: %s", dict(zip(classes, counts, strict=True)))
+
+    model_cfg = dict(cfg["model"])
+    model_cfg.setdefault("num_classes", len(classes))
+    model = build_model(model_cfg).to(device)
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    logger.info(
+        "模型 %s | 参数量 %.2fM（可训练 %.2fM）",
+        model_cfg.get("backbone", {}).get("arch", model_cfg["name"]),
+        sum(p.numel() for p in model.parameters()) / 1e6,
+        trainable / 1e6,
+    )
+
+    trainer = ClassificationTrainer(
+        model=model,
+        train_loader=loaders["train"],
+        val_loader=loaders["val"],
+        cfg=cfg.get("train", {}),
+        device=device,
+        work_dir=work_dir,
+        classes=classes,
+        train_counts=counts,
+        full_config=cfg.to_dict(),
+    )
+
+    if args.eval_only:
+        checkpoint = args.checkpoint or (work_dir / "best.pth")
+        if not Path(checkpoint).exists():
+            logger.error("找不到权重 %s，先训练或用 --checkpoint 指定", checkpoint)
+            return 1
+        from engine import load_checkpoint
+
+        load_checkpoint(checkpoint, model, map_location=device)
+        metrics = trainer.evaluate()
+        logger.info("评估 %s", checkpoint)
+        for key, value in metrics.items():
+            logger.info("  %-24s %.4f", key, value)
+        log_per_class(logger, trainer)
+        write_report(Path(args.report or work_dir / "eval_report.json"), trainer, metrics)
+        return 0
+
+    if args.resume:
+        trainer.resume(args.resume)
+
+    cfg.dump(work_dir / "config.yaml")
+    best = trainer.fit()
+    logger.info("训练结束，最佳 %s=%.4f", trainer.monitor, best)
+
+    # 用 best.pth 再评一次，产出逐类明细与混淆矩阵
+    from engine import load_checkpoint
+
+    best_path = work_dir / "best.pth"
+    if best_path.exists():
+        load_checkpoint(best_path, model, map_location=device)
+        metrics = trainer.evaluate()
+        logger.info("best.pth 在验证集上的表现:")
+        for key, value in metrics.items():
+            logger.info("  %-24s %.4f", key, value)
+        log_per_class(logger, trainer)
+        write_report(Path(args.report or work_dir / "eval_report.json"), trainer, metrics)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

configs/classify/film.yaml

@@ -0,0 +1,78 @@
+# 在人工归类好的数据集上训练缺陷分类模型。
+#
+# 数据来自 scripts/export_clusters.py：自监督聚类 + 人工整理后的 15 个类别。
+#
+# 用法：
+#   uv run python classify.py --config configs/classify/film.yaml
+#   uv run python classify.py --config configs/classify/film.yaml --eval-only \
+#       --checkpoint runs/classify/best.pth
+
+work_dir: runs/classify
+seed: 42
+device: auto
+
+data:
+  root: data/orgnized_from_r101_k40
+  val_ratio: 0.2
+  exclude: []           # 需要剔除某些类别时写在这里，如 [unknown, others]
+  classes: null         # null = 按目录名排序自动确定；训练与评估必须一致
+
+  # 预处理沿用聚类阶段的结论：
+  #   - 拉伸到 128x128 而不是 letterbox——尺寸与宽高比是 ROI 裁剪的产物，不是缺陷属性；
+  #   - 逐图标准化抹掉整体亮度/对比度。这里尤其重要：15 个类别里有几类几乎只来自
+  #     单一原始数据集（ripple 几乎全是 pearlescent、grease 与 zigzag 全是 film），
+  #     不抹掉亮度的话模型可能靠"认出这是哪个数据集"就把类别猜对，看着准确率很高，
+  #     换个批次的数据就崩。
+  image_size: [128, 128]
+  train_transforms:
+    - {name: Resize, size: [128, 128]}
+    - {name: RandomHorizontalFlip, p: 0.5}
+    - {name: RandomVerticalFlip, p: 0.5}
+    - {name: ToTensor}
+    - {name: PerImageStandardize}
+  val_transforms:
+    - {name: Resize, size: [128, 128]}
+    - {name: ToTensor}
+    - {name: PerImageStandardize}
+
+  loader:
+    batch_size: 128
+    num_workers: 8
+
+model:
+  name: classifier
+  backbone:
+    name: torchvision
+    arch: resnet18
+    pretrained: true
+    in_channels: 3
+  classification:
+    dropout: 0.2
+  # 用自监督预训练的骨架初始化——那正是前面聚类阶段的产物。
+  # 例：runs/simclr_r101/simclr_backbone.pth（注意 backbone 配置要与自监督时一致）
+  init_from: null
+  frozen_backbone: false      # true = 冻结骨架只训分类头（线性探针）
+
+train:
+  epochs: 30
+  amp: auto
+  clip_grad: 5.0
+  log_interval: 20
+
+  # 主指标用 macro-F1 而不是准确率：数据从 normal 2679 张到 others 65 张差 41 倍，
+  # 全预测成 normal 也有 24% 准确率，但 macro-F1 会很低。
+  monitor: macro_f1
+
+  # 按类别频率加权，避免小类被淹没；不需要时设 none。
+  class_weight: balanced
+  label_smoothing: 0.05
+
+  optimizer: {name: adamw, lr: 0.0003, weight_decay: 0.0001}
+  scheduler: {name: cosine}
+
+  swanlab:
+    enabled: true
+    project: film-classify
+    name: null
+    mode: local
+    tags: [classify, resnet18]

datasets/__init__.py

@@ -1,6 +1,7 @@
 """数据集：抛出样本的 Dataset 类与相应的数据变换、标签编码。"""
 
 from datasets.build import build_dataloader, build_dataset
+from datasets.classification import FolderClassificationDataset, classification_collate
 from datasets.film_defect import FilmDefectDataset, collate_fn
 from datasets.registry import DATASETS
 from datasets.target_encoder import CenterTargetEncoder
@@ -10,6 +11,8 @@ from datasets.yolo import format_label_lines, load_classes, parse_label_file, wr
 
 __all__ = [
     "build_dataset",
+    "FolderClassificationDataset",
+    "classification_collate",
     "build_dataloader",
     "FilmDefectDataset",
     "collate_fn",

datasets/classification.py

@@ -0,0 +1,122 @@
+"""按目录组织的分类数据集：`root/<类别名>/*.png`，子目录名即标签。
+
+这是 scripts/export_clusters.py 的产物结构——聚类 + 人工归类之后得到的数据集。
+没有预先划分 train/val，这里按类别分层、用固定随机种子切分，保证两次运行拿到同一份划分。
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from pathlib import Path
+from typing import Any
+
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+
+from datasets.registry import DATASETS
+from datasets.transforms import build_transforms
+from datasets.yolo import IMAGE_SUFFIXES
+
+SPLITS = ("train", "val", "all")
+
+
+@DATASETS.register("folder_classification")
+class FolderClassificationDataset(Dataset):
+    """Args:
+    root: 数据集根目录，其下每个子目录是一个类别。
+    split: "train" / "val" / "all"。
+    val_ratio: 验证集比例，**按类别分别切分**——小类只有几十张，随机整体切分
+        很容易让某个类在验证集里一张都没有。
+    seed: 划分用的随机种子；同一个 seed 两次运行得到完全相同的划分。
+    classes: 显式指定类别顺序；缺省按目录名排序。训练与评估必须一致，
+        否则标签下标对不上。
+    transforms: transform 配置列表，见 `datasets.transforms.build_transforms`。
+    exclude: 要跳过的子目录名。
+    """
+
+    def __init__(
+        self,
+        root: str | Path,
+        split: str = "train",
+        val_ratio: float = 0.2,
+        seed: int = 42,
+        classes: Sequence[str] | None = None,
+        transforms: Sequence[Mapping[str, Any]] = (),
+        gray: bool = False,
+        exclude: Sequence[str] = (),
+    ) -> None:
+        if split not in SPLITS:
+            raise ValueError(f"未知 split: {split}，可选 {SPLITS}")
+        self.root = Path(root)
+        if not self.root.is_dir():
+            raise FileNotFoundError(f"数据集目录不存在: {self.root}")
+
+        self.split = split
+        self.gray = gray
+        self.transforms = build_transforms(transforms)
+
+        excluded = set(exclude)
+        found = sorted(
+            p.name for p in self.root.iterdir() if p.is_dir() and p.name not in excluded
+        )
+        self.classes = list(classes) if classes else found
+        if not self.classes:
+            raise FileNotFoundError(f"{self.root} 下没有找到类别子目录")
+        missing = [c for c in self.classes if not (self.root / c).is_dir()]
+        if missing:
+            raise FileNotFoundError(f"以下类别在 {self.root} 下不存在: {missing}")
+        self.class_to_idx = {name: i for i, name in enumerate(self.classes)}
+
+        self.samples = self._collect(val_ratio, seed)
+        if not self.samples:
+            raise ValueError(f"{self.root} 的 {split} 划分里没有样本")
+
+    def _collect(self, val_ratio: float, seed: int) -> list[tuple[Path, int]]:
+        samples: list[tuple[Path, int]] = []
+        for label, name in enumerate(self.classes):
+            files = sorted(
+                p
+                for p in (self.root / name).rglob("*")
+                if p.is_file() and p.suffix.lower() in IMAGE_SUFFIXES
+            )
+            if not files:
+                continue
+            # 逐类切分：小类只有几十张，整体随机切很容易让它在验证集里挂零
+            generator = torch.Generator().manual_seed(seed + label)
+            order = torch.randperm(len(files), generator=generator).tolist()
+            n_val = max(1, round(len(files) * val_ratio)) if len(files) > 1 else 0
+            val_idx = set(order[:n_val])
+            for index, path in enumerate(files):
+                in_val = index in val_idx
+                if self.split == "all" or (self.split == "val") == in_val:
+                    samples.append((path, label))
+        return samples
+
+    def class_counts(self) -> list[int]:
+        counts = [0] * len(self.classes)
+        for _, label in self.samples:
+            counts[label] += 1
+        return counts
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, index: int) -> dict:
+        path, label = self.samples[index]
+        with Image.open(path) as raw:
+            image = raw.convert("L" if self.gray else "RGB")
+        sample = self.transforms({"image": image})
+        return {
+            "image": sample["image"],
+            "label": label,
+            "meta": {"path": str(path), "index": index, "class": self.classes[label]},
+        }
+
+
+def classification_collate(batch: list[dict]) -> dict:
+    return {
+        "images": torch.stack([item["image"] for item in batch]),
+        "labels": torch.tensor([item["label"] for item in batch], dtype=torch.long),
+        "meta": [item["meta"] for item in batch],
+    }

engine/__init__.py

@@ -1,6 +1,11 @@
 """训练 / 评估流程。"""
 
 from engine.checkpoint import load_checkpoint, save_checkpoint
+from engine.classification import (
+    ClassificationMetric,
+    ClassificationTrainer,
+    class_weights,
+)
 from engine.evaluator import Evaluator
 from engine.losses import DetectionLoss
 from engine.optim import build_optimizer, build_scheduler
@@ -8,6 +13,9 @@ from engine.trainer import Trainer
 
 __all__ = [
     "Trainer",
+    "ClassificationTrainer",
+    "ClassificationMetric",
+    "class_weights",
     "Evaluator",
     "DetectionLoss",
     "build_optimizer",

engine/classification.py

@@ -0,0 +1,288 @@
+"""分类任务的训练、评估与指标。
+
+数据极不均衡（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

models/__init__.py

@@ -5,6 +5,7 @@
 
 from models import backbones, heads, necks  # noqa: F401  触发注册
 from models.build import build_model
+from models.classifier import DefectClassifier
 from models.defect_net import DefectNet
 from models.postprocess import decode_heatmap
 from models.registry import BACKBONES, HEADS, MODELS, NECKS
@@ -12,6 +13,7 @@ from models.registry import BACKBONES, HEADS, MODELS, NECKS
 __all__ = [
     "build_model",
     "DefectNet",
+    "DefectClassifier",
     "decode_heatmap",
     "BACKBONES",
     "NECKS",

models/classifier.py

@@ -0,0 +1,82 @@
+"""缺陷分类模型：卷积骨架 + 分类头。
+
+与 `DefectNet` 共用同一套 backbone 注册表，所以自监督预训练出来的骨架
+（`runs/*/simclr_backbone.pth`）可以直接拿来初始化——那正是前面聚类阶段的产物。
+按项目约定只用 convolution，不含 attention / transformer。
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+import torch
+from torch import nn
+
+from models.heads.classification import ClassificationHead
+from models.registry import BACKBONES, MODELS
+from utils.logging import get_logger
+
+logger = get_logger("models.classifier")
+
+
+@MODELS.register("classifier")
+class DefectClassifier(nn.Module):
+    """Args:
+    num_classes: 类别数。
+    backbone: backbone 注册表的配置。
+    classification: 分类头参数（`hidden`/`dropout` 等）。
+    init_from: 自监督骨架权重路径；只加载 backbone，分类头仍随机初始化。
+    frozen_backbone: 冻结骨架只训分类头（线性探针，用来单独衡量自监督特征的好坏）。
+    """
+
+    def __init__(
+        self,
+        num_classes: int,
+        backbone: Mapping[str, Any] | None = None,
+        classification: Mapping[str, Any] | None = None,
+        init_from: str | Path | None = None,
+        frozen_backbone: bool = False,
+    ) -> None:
+        super().__init__()
+        self.num_classes = num_classes
+        self.backbone = BACKBONES.build(dict(backbone or {"name": "tiny_conv"}))
+        self.head = ClassificationHead(
+            self.backbone.out_channels, num_classes, **dict(classification or {})
+        )
+
+        if init_from:
+            self.load_backbone(init_from)
+        if frozen_backbone:
+            for param in self.backbone.parameters():
+                param.requires_grad_(False)
+            logger.info("骨架已冻结，只训练分类头（线性探针）")
+
+    def load_backbone(self, path: str | Path) -> None:
+        """加载自监督预训练的骨架权重。
+
+        只取 backbone 那部分：simclr_backbone.pth 里存的就是骨架 state_dict，
+        分类头是新的、必须从头学。
+        """
+        path = Path(path)
+        if not path.exists():
+            raise FileNotFoundError(f"预训练骨架不存在: {path}")
+        state = torch.load(path, map_location="cpu", weights_only=False)
+        weights = state.get("model", state.get("backbone", state))
+        missing, unexpected = self.backbone.load_state_dict(weights, strict=False)
+        logger.info(
+            "已从 %s 加载骨架权重（缺失 %d 项，多余 %d 项，已训 %s 轮）",
+            path,
+            len(missing),
+            len(unexpected),
+            state.get("epochs", "?"),
+        )
+        if len(missing) > len(list(self.backbone.state_dict())) // 2:
+            raise ValueError(
+                f"{path} 与当前骨架结构对不上（缺失 {len(missing)} 项）；"
+                "确认 backbone 配置与自监督训练时一致"
+            )
+
+    def forward(self, images: torch.Tensor) -> torch.Tensor:
+        return self.head(self.backbone(images))

tests/test_classification.py

@@ -0,0 +1,274 @@
+"""分类数据集、模型与训练引擎。"""
+
+import sys
+
+import pytest
+
+torch = pytest.importorskip("torch")
+pytest.importorskip("PIL")
+
+import numpy as np  # noqa: E402
+from PIL import Image  # noqa: E402
+
+from datasets import FolderClassificationDataset, classification_collate  # noqa: E402
+from engine.classification import (  # noqa: E402
+    ClassificationMetric,
+    ClassificationTrainer,
+    class_weights,
+)
+from models import build_model  # noqa: E402
+
+TRANSFORMS = [{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}]
+
+
+@pytest.fixture
+def dataset_root(tmp_path):
+    """三个类别，样本数刻意不均衡（20 / 6 / 2）。"""
+    rng = np.random.default_rng(0)
+    root = tmp_path / "dataset"          # 单独一层：输出目录不能落在数据集根下，
+    root.mkdir()                         # 否则会被当成一个新类别扫进来
+    for name, count in (("normal", 20), ("grease", 6), ("others", 2)):
+        directory = root / name
+        directory.mkdir()
+        for i in range(count):
+            Image.fromarray(rng.integers(0, 255, (18, 24, 3), dtype=np.uint8)).save(
+                directory / f"{i}.png"
+            )
+    return root
+
+
+def test_classes_come_from_subdirectory_names(dataset_root):
+    dataset = FolderClassificationDataset(dataset_root, split="all", transforms=TRANSFORMS)
+    assert dataset.classes == ["grease", "normal", "others"]  # 按目录名排序
+    assert len(dataset) == 28
+    assert dataset.class_counts() == [6, 20, 2]
+
+
+def test_split_is_stratified_and_deterministic(dataset_root):
+    train = FolderClassificationDataset(dataset_root, split="train", transforms=TRANSFORMS)
+    val = FolderClassificationDataset(dataset_root, split="val", transforms=TRANSFORMS)
+
+    assert len(train) + len(val) == 28
+    # 逐类切分：即使只有 2 张的类别也要在验证集里留下样本
+    assert all(count > 0 for count in val.class_counts()), val.class_counts()
+    assert all(count > 0 for count in train.class_counts())
+
+    again = FolderClassificationDataset(dataset_root, split="val", transforms=TRANSFORMS)
+    assert [p for p, _ in val.samples] == [p for p, _ in again.samples]
+
+    # 训练与验证不能有交集
+    assert not {p for p, _ in train.samples} & {p for p, _ in val.samples}
+
+
+def test_different_seed_changes_the_split(dataset_root):
+    a = FolderClassificationDataset(dataset_root, split="val", seed=1, transforms=TRANSFORMS)
+    b = FolderClassificationDataset(dataset_root, split="val", seed=2, transforms=TRANSFORMS)
+    assert [p for p, _ in a.samples] != [p for p, _ in b.samples]
+
+
+def test_exclude_and_explicit_classes(dataset_root):
+    dataset = FolderClassificationDataset(
+        dataset_root, split="all", transforms=TRANSFORMS, exclude=["others"]
+    )
+    assert dataset.classes == ["grease", "normal"]
+
+    ordered = FolderClassificationDataset(
+        dataset_root, split="all", transforms=TRANSFORMS, classes=["others", "normal"]
+    )
+    assert ordered.class_to_idx == {"others": 0, "normal": 1}
+
+    with pytest.raises(FileNotFoundError, match="不存在"):
+        FolderClassificationDataset(dataset_root, classes=["nope"], transforms=TRANSFORMS)
+
+
+def test_sample_and_collate_shapes(dataset_root):
+    dataset = FolderClassificationDataset(dataset_root, split="all", transforms=TRANSFORMS)
+    sample = dataset[0]
+    assert sample["image"].shape == (3, 32, 32)
+    assert isinstance(sample["label"], int)
+    assert sample["meta"]["class"] == dataset.classes[sample["label"]]
+
+    batch = classification_collate([dataset[0], dataset[1]])
+    assert batch["images"].shape == (2, 3, 32, 32)
+    assert batch["labels"].shape == (2,) and batch["labels"].dtype == torch.long
+
+
+def test_class_weights_favour_rare_classes():
+    weights = class_weights([100, 10, 1])
+    assert weights[2] > weights[1] > weights[0]
+    assert float(weights.mean()) == pytest.approx(1.0)
+    assert class_weights([1, 2], mode="none") is None
+    with pytest.raises(ValueError, match="未知 class_weight"):
+        class_weights([1, 2], mode="sqrt")
+
+
+def test_metric_matches_hand_computed_values():
+    metric = ClassificationMetric(["a", "b"])
+    # 真实 a,a,b,b -> 预测 a,b,b,b：a 召回 0.5，b 精确 2/3
+    metric.update(torch.tensor([[9.0, 0.0], [0.0, 9.0], [0.0, 9.0], [0.0, 9.0]]),
+                  torch.tensor([0, 0, 1, 1]))
+    report = metric.per_class()
+    assert report["a"] == pytest.approx({"precision": 1.0, "recall": 0.5,
+                                         "f1": 2 / 3, "support": 2})
+    assert report["b"]["precision"] == pytest.approx(2 / 3)
+
+    metrics = metric.compute()
+    assert metrics["accuracy"] == pytest.approx(0.75)
+    assert metrics["macro_f1"] == pytest.approx((2 / 3 + 0.8) / 2)
+
+
+def test_macro_f1_punishes_majority_only_prediction():
+    """全预测成多数类：准确率还行，但 macro-F1 很低——这正是选它当主指标的原因。"""
+    metric = ClassificationMetric(["major", "minor"])
+    logits = torch.tensor([[9.0, 0.0]]).repeat(100, 1)
+    targets = torch.tensor([0] * 90 + [1] * 10)
+    metric.update(logits, targets)
+
+    metrics = metric.compute()
+    assert metrics["accuracy"] == pytest.approx(0.9)
+    assert metrics["macro_f1"] < 0.5
+    assert metrics["balanced_accuracy"] == pytest.approx(0.5)
+
+
+def test_classifier_is_convolution_only_and_shapes_match():
+    model = build_model({
+        "name": "classifier", "num_classes": 4,
+        "backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]},
+        "classification": {"dropout": 0.1},
+    }).eval()
+    with torch.no_grad():
+        assert model(torch.randn(3, 3, 64, 64)).shape == (3, 4)
+    forbidden = (torch.nn.MultiheadAttention, torch.nn.Transformer, torch.nn.LSTM)
+    assert not any(isinstance(m, forbidden) for m in model.modules())
+
+
+def test_init_from_loads_self_supervised_backbone(tmp_path):
+    """自监督骨架能直接拿来初始化分类模型——这是聚类阶段产物的下游用途。"""
+    from models.registry import BACKBONES
+
+    backbone_cfg = {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}
+    pretrained = BACKBONES.build(dict(backbone_cfg))
+    with torch.no_grad():
+        for param in pretrained.parameters():
+            param.fill_(0.5)
+    path = tmp_path / "simclr_backbone.pth"
+    torch.save({"model": pretrained.state_dict(), "epochs": 40}, path)
+
+    model = build_model({"name": "classifier", "num_classes": 3,
+                         "backbone": dict(backbone_cfg), "init_from": str(path)})
+    for param in model.backbone.parameters():
+        assert torch.allclose(param, torch.full_like(param, 0.5))
+
+    with pytest.raises(FileNotFoundError):
+        build_model({"name": "classifier", "num_classes": 3,
+                     "backbone": dict(backbone_cfg), "init_from": str(tmp_path / "nope.pth")})
+
+
+def test_frozen_backbone_only_trains_the_head():
+    model = build_model({
+        "name": "classifier", "num_classes": 3,
+        "backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]},
+        "frozen_backbone": True,
+    })
+    assert not any(p.requires_grad for p in model.backbone.parameters())
+    assert all(p.requires_grad for p in model.head.parameters())
+
+
+def _trainer(dataset_root, tmp_path, **cfg_extra):
+    from torch.utils.data import DataLoader
+
+    datasets = {
+        split: FolderClassificationDataset(dataset_root, split=split, transforms=TRANSFORMS)
+        for split in ("train", "val")
+    }
+    loaders = {
+        split: DataLoader(ds, batch_size=4, shuffle=(split == "train"),
+                          collate_fn=classification_collate)
+        for split, ds in datasets.items()
+    }
+    model = build_model({"name": "classifier", "num_classes": len(datasets["train"].classes),
+                         "backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}})
+    cfg = {"epochs": 2, "amp": False, "log_interval": 999, "scheduler": {"name": "cosine"},
+           "optimizer": {"name": "adamw", "lr": 1e-3}, "swanlab": {"enabled": False}, **cfg_extra}
+    return ClassificationTrainer(
+        model, loaders["train"], loaders["val"], cfg, torch.device("cpu"), tmp_path / "run",
+        classes=datasets["train"].classes, train_counts=datasets["train"].class_counts(),
+    )
+
+
+def test_training_runs_and_saves_best_by_macro_f1(dataset_root, tmp_path):
+    trainer = _trainer(dataset_root, tmp_path)
+    best = trainer.fit()
+
+    assert (tmp_path / "run" / "best.pth").exists()
+    assert (tmp_path / "run" / "last.pth").exists()
+    assert best > float("-inf")
+    assert trainer.monitor == "macro_f1"   # 不均衡数据下的默认主指标
+
+    metrics = trainer.evaluate()
+    assert {"accuracy", "macro_f1", "weighted_f1", "balanced_accuracy", "loss"} <= set(metrics)
+
+
+def test_resume_takes_hyperparams_from_config(dataset_root, tmp_path):
+    """与检测/自监督一致：检查点只提供进度，超参以配置为准。"""
+    trainer = _trainer(dataset_root, tmp_path, epochs=2,
+                       optimizer={"name": "adamw", "lr": 1e-3, "weight_decay": 1e-4})
+    trainer.fit()
+
+    resumed = _trainer(dataset_root, tmp_path, epochs=6,
+                       optimizer={"name": "adamw", "lr": 5e-4, "weight_decay": 0.05})
+    resumed.resume(tmp_path / "run" / "last.pth")
+
+    assert resumed.start_epoch == 2
+    groups = resumed.optimizer.param_groups
+    assert groups[0]["initial_lr"] == pytest.approx(5e-4)
+    assert groups[0]["weight_decay"] == pytest.approx(0.05)
+    assert resumed.scheduler.T_max == 6
+
+
+def test_cli_train_and_eval_only(dataset_root, tmp_path, monkeypatch):
+    import classify
+
+    config = tmp_path / "cfg.yaml"
+    config.write_text(
+        f"""
+work_dir: {tmp_path / "run_cli"}
+seed: 42
+device: cpu
+data:
+  root: {dataset_root}
+  val_ratio: 0.25
+  exclude: []
+  classes: null
+  train_transforms: [{{name: Resize, size: [32, 32]}}, {{name: ToTensor}}]
+  val_transforms: [{{name: Resize, size: [32, 32]}}, {{name: ToTensor}}]
+  loader: {{batch_size: 4, num_workers: 0}}
+model:
+  name: classifier
+  backbone: {{name: tiny_conv, widths: [8, 16], depths: [1, 1]}}
+  classification: {{dropout: 0.0}}
+  init_from: null
+  frozen_backbone: false
+train:
+  epochs: 1
+  amp: false
+  log_interval: 999
+  monitor: macro_f1
+  class_weight: balanced
+  optimizer: {{name: adamw, lr: 0.001}}
+  scheduler: {{name: cosine}}
+  swanlab: {{enabled: false}}
+""",
+        encoding="utf-8",
+    )
+
+    monkeypatch.setattr(sys, "argv", ["classify.py", "--config", str(config)])
+    assert classify.main() == 0
+    assert (tmp_path / "run_cli" / "eval_report.json").exists()
+
+    monkeypatch.setattr(
+        sys, "argv",
+        ["classify.py", "--config", str(config), "--eval-only",
+         "--checkpoint", str(tmp_path / "run_cli" / "best.pth")],
+    )
+    assert classify.main() == 0