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>
2026-09-02 17:55:20 +08:00
parent ef64e5bf77
commit b0d0dfffbb
9 changed files with 344 additions and 0 deletions
@@ -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())