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>
This commit is contained in:
co-authored by
Claude Opus 5
parent
ef64e5bf77
commit
b0d0dfffbb
@@ -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())
|
||||
Reference in New Issue
Block a user