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>
2026-09-10 14:58:41 +08:00
parent 33d28dcc7c
commit 33fdf396ee
10 changed files with 1085 additions and 1 deletions
@@ -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