#!/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 datasets.tile_classification import YoloTileClassificationDataset 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")): if split_cfg.get("format", "folder") == "yolo_tiles": datasets[split] = YoloTileClassificationDataset( yaml_path=split_cfg["yaml"], split=split, image_size=split_cfg["image_size"], transforms=split_cfg[transforms_key], classes=classes or split_cfg.get("classes"), ) continue 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())