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>
275 lines
11 KiB
Python
275 lines
11 KiB
Python
"""分类数据集、模型与训练引擎。"""
|
||
|
||
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
|