ModelTrainEval/classify.py

#!/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())