189 lines
7.0 KiB
Python
Executable File
189 lines
7.0 KiB
Python
Executable File
#!/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())
|