feat: 新增缺陷分类训练与评估链路
在聚类 + 人工归类得到的数据集（data/<集>/<类别名>/*.png）上训练分类模型。 - datasets/classification.py：按目录组织的分类数据集，**逐类分层**切分 train/val （others 只有 65 张，整体随机切很容易让它在验证集里挂零），固定种子可复现 - models/classifier.py：卷积骨架 + 分类头，与 DefectNet 共用 backbone 注册表， 因此自监督产出的 simclr_backbone.pth 可用 init_from 直接初始化； frozen_backbone 可做线性探针 - engine/classification.py：训练循环 + 指标。主指标用 macro-F1 而非准确率—— 数据从 normal 2679 张到 others 65 张差 41 倍，全预测成 normal 也有 24% 准确率； 损失按类别频率加权，避免小类被淹没；输出逐类 P/R/F1 与混淆矩阵 - classify.py：训练/评估入口，--eval-only 可单独评估并导出 eval_report.json - configs/classify/film.yaml：预处理沿用聚类阶段的结论（拉伸 + 逐图标准化）。 抹掉亮度在这里尤其重要：有几类几乎只来自单一原始数据集，不抹掉的话模型可能 靠"认出这是哪个数据集"作弊 - 续训语义与检测/自监督一致：检查点只提供进度，超参以配置为准 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
33d28dcc7c
commit
33fdf396ee
Executable
+180
@@ -0,0 +1,180 @@
|
||||
#!/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 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")):
|
||||
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())
|
||||
Reference in New Issue
Block a user