Shaun ZhouandClaude Opus 5
141cb907f2
移除"训练必须有 Nvidia GPU"的硬性守卫，改为按后端能力分派。
- utils/device: resolve_device("auto") 按 cuda > mps > cpu 挑最快的可用后端；
显式指定不可用的后端时直接报错，不会悄悄退回 CPU——训练慢几十倍却不吭声是最坏的情况
- AMP 随后端取默认值：CUDA 开、MPS 关。MPS 上 autocast 数值正常但实测更慢
（tiny_conv @128²：fp32 76 ms/步、fp16 90 ms、bf16 108 ms），需要时 amp: true 可强开
- PeakMemory 统一显存峰值统计：CUDA 读内置计数器，MPS 无峰值 API 改用逐 step 采样
（是下界，日志与产出里已标注）
- profile_inference 支持 MPS，并提示 torch.profiler 没有 MPS 的 ProfilerActivity，
算子明细只反映 CPU 侧下发耗时
- train.py / eval.py 新增 --device，配置默认改为 auto
- check_env 改为列出全部可用后端与 auto 的选择结果
实测：tiny_conv 训练在 MPS 上比 CPU 快约 32 倍。
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
89 lines
3.2 KiB
Python
Executable File
89 lines
3.2 KiB
Python
Executable File
#!/usr/bin/env python
|
||
"""评估入口。
|
||
|
||
用法:
|
||
python eval.py --config configs/eval/baseline.yaml --checkpoint runs/baseline/best.pth
|
||
|
||
设备同样支持 auto / cuda / mps / cpu；部署前回归可以显式 `--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 DeviceNotAvailableError, describe_device, 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("--device", default=None, help="覆盖配置中的 device：auto / cuda / mps / cpu")
|
||
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
|
||
|
||
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))
|
||
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())
|