feat: 训练与评估入口及 yaml 配置

- train.py: 先过 require_cuda() 守卫再建模型/数据，支持 --resume 与 --set 覆盖
- eval.py: 允许 CPU，输出定位与分类指标并可写 json
- configs: model/dataset/train/eval 四类，通过 `_base_` 组合；
  classes 既可直接写列表，也可指向 YOLO data.yaml 由框架读 names

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

configs/dataset/film_defect.yaml

@@ -0,0 +1,51 @@
+# 薄膜缺陷数据集（YOLO 风格标注：labels/*.txt，每行 `class_id cx cy w h`，归一化坐标）。
+# 类别顺序即标注里的 class_id，改动顺序等于改变标签含义，必须重新训练。
+# 也可以写成 data.yaml 的路径，从其中的 names 读类别：classes: data/film/data.yaml
+classes: [scratch, bubble, particle, crease, stain]
+
+# 图像级标签形式；模型的分类头、分类损失、分类指标都跟随它。
+#   multilabel（多类）：一张图可同时含多种缺陷，多热标签 + BCE，全 0 表示 OK 片；
+#   multiclass（单类）：一张图只归一个类别，互斥 softmax + CE，
+#                       分类头多一维表示"无缺陷"（下标 = len(classes)），
+#                       一图出现多种缺陷时取框数最多的类别。
+# 无论取哪种，定位分支始终按 classes 逐类输出 heatmap。
+label_mode: multilabel
+
+image_size: [512, 512]
+gray: false
+
+data:
+  train:
+    dataset:
+      name: film_defect
+      img_dir: data/film/images/train
+      label_dir: null          # 缺省按 YOLO 约定取 data/film/labels/train
+      gray: false
+      max_objects: 64
+      transforms:
+        - {name: Resize, size: [512, 512]}
+        - {name: RandomHorizontalFlip, p: 0.5}
+        - {name: RandomVerticalFlip, p: 0.5}
+        - {name: ColorJitter, brightness: 0.2, contrast: 0.2}
+        - {name: ToTensor}
+        - {name: Normalize, mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225]}
+    loader:
+      batch_size: 16
+      shuffle: true
+      num_workers: 8
+      drop_last: true
+
+  val:
+    dataset:
+      name: film_defect
+      img_dir: data/film/images/val
+      label_dir: null
+      gray: false
+      transforms:
+        - {name: Resize, size: [512, 512]}
+        - {name: ToTensor}
+        - {name: Normalize, mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225]}
+    loader:
+      batch_size: 16
+      shuffle: false
+      num_workers: 4

configs/eval/baseline.yaml

@@ -0,0 +1,12 @@
+_base_:
+  - ../train/baseline.yaml
+
+# 评估允许在 CPU 上跑（部署前的精度回归），训练则必须用 GPU。
+device: cuda
+checkpoint: runs/baseline/best.pth
+
+eval:
+  split: val
+  score_threshold: 0.3
+  iou_threshold: 0.5
+  topk: 50

configs/model/resnet18.yaml

@@ -0,0 +1,15 @@
+# torchvision resnet18：stride 32 + 3 次上采样 => 输出 stride 4。
+model:
+  name: defect_net
+  backbone:
+    name: torchvision
+    arch: resnet18
+    pretrained: true
+    in_channels: 3
+  neck:
+    name: upsample
+    channels: 64
+    num_upsamples: 3
+  head:
+    hidden_channels: 64
+  with_classification: true

configs/model/tiny_conv.yaml

@@ -0,0 +1,18 @@
+# 自研轻量骨架：backbone stride 16 + 2 次上采样 => 输出 stride 4。
+model:
+  name: defect_net
+  backbone:
+    name: tiny_conv
+    in_channels: 3
+    widths: [32, 64, 128, 256]
+    depths: [1, 2, 2, 2]
+    block: dw          # dw=深度可分离(更快) / res=残差(更强)
+  neck:
+    name: upsample
+    channels: 64
+    num_upsamples: 2
+  head:
+    hidden_channels: 64
+  with_classification: true
+  classification:
+    dropout: 0.1

configs/train/baseline.yaml

@@ -0,0 +1,42 @@
+_base_:
+  - ../model/tiny_conv.yaml
+  - ../dataset/film_defect.yaml
+
+work_dir: runs/baseline
+seed: 42
+deterministic: false
+device: cuda            # 训练只允许 Nvidia GPU，见 utils/device.py
+
+train:
+  epochs: 100
+  amp: true
+  clip_grad: 5.0
+  log_interval: 20
+  monitor: f1           # 用于挑选 best.pth 的验证指标
+
+  optimizer:
+    name: adamw
+    lr: 0.001
+    weight_decay: 0.0001
+
+  scheduler:
+    name: cosine
+
+  loss:
+    heatmap_weight: 1.0
+    wh_weight: 0.1
+    offset_weight: 1.0
+    cls_weight: 1.0
+
+  eval:
+    score_threshold: 0.3
+    iou_threshold: 0.5
+    topk: 50
+
+  # 训练过程可视化（SwanLab）。未安装或 enabled=false 时自动跳过。
+  swanlab:
+    enabled: true
+    project: film-defect
+    name: null            # 缺省用 work_dir 的目录名
+    mode: online          # online(云端) | local(本地看板) | offline(仅存盘) | disabled
+    tags: [tiny_conv, center-head]

configs/train/multiclass.yaml

@@ -0,0 +1,12 @@
+_base_:
+  - ./baseline.yaml
+
+# 单类（互斥）分类示例：一张图只归一个类别，分类头多一维表示"无缺陷"。
+label_mode: multiclass
+
+work_dir: runs/multiclass
+train:
+  monitor: cls_f1         # 单类模式下常用分类指标挑选 best.pth
+  swanlab:
+    name: multiclass
+    tags: [tiny_conv, multiclass]

configs/train/resnet18.yaml

@@ -0,0 +1,8 @@
+_base_:
+  - ./baseline.yaml
+  - ../model/resnet18.yaml
+
+work_dir: runs/resnet18
+train:
+  optimizer:
+    lr: 0.0005

eval.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python
+"""评估入口。
+
+用法:
+    python eval.py --config configs/eval/baseline.yaml --checkpoint runs/baseline/best.pth
+
+与训练不同，评估允许在 CPU 上运行（`--set device=cpu`），便于部署前回归。
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+from datasets import build_dataloader, build_dataset, load_classes
+from engine import Evaluator
+from models import build_model
+from utils import check_label_mode, load_config, setup_logging
+from utils.device import resolve_device
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="薄膜缺陷模型评估")
+    parser.add_argument("--config", required=True)
+    parser.add_argument("--checkpoint", default=None, help="覆盖配置中的 checkpoint")
+    parser.add_argument("--split", default=None, help="评估哪个 data 划分，默认取配置的 eval.split")
+    parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE")
+    parser.add_argument("--out", default=None, help="把指标写入 json 文件")
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    cfg = load_config(args.config, args.set)
+    logger = setup_logging()
+
+    eval_cfg = dict(cfg.get("eval", {}))
+    split = args.split or eval_cfg.pop("split", "val")
+    eval_cfg.pop("split", None)
+
+    checkpoint = args.checkpoint or cfg.get("checkpoint")
+    if not checkpoint:
+        logger.error("必须通过 --checkpoint 或配置中的 checkpoint 指定权重")
+        return 1
+
+    device = resolve_device(str(cfg.get("device", "cuda")))
+    classes = load_classes(cfg["classes"])
+
+    # label_mode 由 yaml 决定，模型与数据集必须一致，这里统一下发。
+    label_mode = check_label_mode(str(cfg.get("label_mode", "multilabel")))
+    model_cfg = dict(cfg["model"])
+    model_cfg.setdefault("num_classes", len(classes))
+    model_cfg.setdefault("label_mode", label_mode)
+    model = build_model(model_cfg)
+
+    from engine import load_checkpoint
+
+    load_checkpoint(checkpoint, model, map_location=device)
+    model.to(device).eval()
+
+    split_cfg = cfg["data"][split]
+    dataset = build_dataset(
+        split_cfg["dataset"], classes=classes, stride=model.out_stride, label_mode=label_mode
+    )
+    loader = build_dataloader(dataset, **{**dict(split_cfg.get("loader", {})), "shuffle": False})
+    logger.info("评估 %s 集，共 %d 张图，权重 %s", split, len(dataset), checkpoint)
+
+    metrics = Evaluator(model, loader, device, **eval_cfg).evaluate()
+    for key, value in metrics.items():
+        logger.info("%-16s %.4f", key, value)
+
+    if args.out:
+        Path(args.out).parent.mkdir(parents=True, exist_ok=True)
+        Path(args.out).write_text(json.dumps(metrics, indent=2, ensure_ascii=False), encoding="utf-8")
+        logger.info("指标已写入 %s", args.out)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

train.py

@@ -0,0 +1,104 @@
+#!/usr/bin/env python
+"""训练入口。
+
+用法:
+    python train.py --config configs/train/baseline.yaml
+    python train.py --config configs/train/baseline.yaml --set train.epochs=10 train.optimizer.lr=5e-4
+
+硬性约束：没有 Nvidia GPU 时直接退出（见 CLAUDE.MD / utils/device.py）。
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+import torch
+
+from datasets import build_dataloader, build_dataset, load_classes
+from engine import Trainer
+from models import build_model
+from utils import check_label_mode, load_config, seed_everything, setup_logging
+from utils.device import NoNvidiaGpuError, require_cuda
+
+
+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, help="覆盖配置中的 work_dir")
+    parser.add_argument("--resume", default=None, help="从检查点恢复训练")
+    return parser.parse_args()
+
+
+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/default"))
+    work_dir.mkdir(parents=True, exist_ok=True)
+    logger = setup_logging(work_dir / "train.log")
+
+    try:
+        gpus = require_cuda()
+    except NoNvidiaGpuError as exc:
+        logger.error("%s", exc)
+        return 1
+    logger.info("检测到 GPU: %s", ", ".join(gpus))
+
+    seed_everything(int(cfg.get("seed", 42)), bool(cfg.get("deterministic", False)))
+    device = torch.device(cfg.get("device", "cuda"))
+
+    classes = load_classes(cfg["classes"])
+    # label_mode 由 yaml 决定，模型与数据集必须一致，这里统一下发。
+    label_mode = check_label_mode(str(cfg.get("label_mode", "multilabel")))
+    model_cfg = dict(cfg["model"])
+    model_cfg.setdefault("num_classes", len(classes))
+    model_cfg.setdefault("label_mode", label_mode)
+    model = build_model(model_cfg).to(device)
+    logger.info(
+        "模型 %s 输出 stride=%d，参数量 %.2fM",
+        model_cfg["name"],
+        model.out_stride,
+        sum(p.numel() for p in model.parameters()) / 1e6,
+    )
+
+    loaders = {}
+    for split in ("train", "val"):
+        split_cfg = cfg["data"].get(split)
+        if not split_cfg:
+            loaders[split] = None
+            continue
+        dataset = build_dataset(
+            split_cfg["dataset"], classes=classes, stride=model.out_stride, label_mode=label_mode
+        )
+        loaders[split] = build_dataloader(dataset, **dict(split_cfg.get("loader", {})))
+        logger.info("%s 集样本数: %d", split, len(dataset))
+
+    trainer = Trainer(
+        model=model,
+        train_loader=loaders["train"],
+        val_loader=loaders["val"],
+        cfg=cfg.get("train", {}),
+        device=device,
+        work_dir=work_dir,
+        full_config=cfg.to_dict(),
+    )
+
+    if args.resume:
+        from engine import load_checkpoint
+
+        state = load_checkpoint(args.resume, model, trainer.optimizer, map_location=device)
+        trainer.start_epoch = int(state.get("epoch", 0)) + 1
+        trainer.best_metric = float(state.get("best_metric") or float("-inf"))
+        logger.info("已从 %s 恢复，起始 epoch=%d", args.resume, trainer.start_epoch)
+
+    cfg.dump(work_dir / "config.yaml")
+    best = trainer.fit()
+    logger.info("训练结束，最佳 %s=%.4f", trainer.monitor, best)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())