diff --git a/README.md b/README.md index 14d15cc..3308677 100644 --- a/README.md +++ b/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）。 diff --git a/classify.py b/classify.py new file mode 100755 index 0000000..4e9ccdd --- /dev/null +++ b/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()) diff --git a/configs/classify/film.yaml b/configs/classify/film.yaml new file mode 100644 index 0000000..c27abde --- /dev/null +++ b/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] diff --git a/datasets/__init__.py b/datasets/__init__.py index c7896ec..3f70885 100644 --- a/datasets/__init__.py +++ b/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", diff --git a/datasets/classification.py b/datasets/classification.py new file mode 100644 index 0000000..f3a258b --- /dev/null +++ b/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], + } diff --git a/engine/__init__.py b/engine/__init__.py index 8d7d7a2..2c5e4f7 100644 --- a/engine/__init__.py +++ b/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", diff --git a/engine/classification.py b/engine/classification.py new file mode 100644 index 0000000..0b42d51 --- /dev/null +++ b/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 diff --git a/models/__init__.py b/models/__init__.py index 6cab961..e540694 100644 --- a/models/__init__.py +++ b/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", diff --git a/models/classifier.py b/models/classifier.py new file mode 100644 index 0000000..baa2084 --- /dev/null +++ b/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)) diff --git a/tests/test_classification.py b/tests/test_classification.py new file mode 100644 index 0000000..4b61f37 --- /dev/null +++ b/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