diff --git a/configs/dataset/film_defect.yaml b/configs/dataset/film_defect.yaml new file mode 100644 index 0000000..d80453a --- /dev/null +++ b/configs/dataset/film_defect.yaml @@ -0,0 +1,51 @@ +# 薄膜缺陷数据集(YOLO 风格标注:labels/*.txt,每行 `class_id cx cy w h`,归一化坐标)。 +# 类别顺序即标注里的 class_id,改动顺序等于改变标签含义,必须重新训练。 +# 也可以写成 data.yaml 的路径,从其中的 names 读类别:classes: data/film/data.yaml +classes: [scratch, bubble, particle, crease, stain] + +# 图像级标签形式;模型的分类头、分类损失、分类指标都跟随它。 +# multilabel(多类):一张图可同时含多种缺陷,多热标签 + BCE,全 0 表示 OK 片; +# multiclass(单类):一张图只归一个类别,互斥 softmax + CE, +# 分类头多一维表示"无缺陷"(下标 = len(classes)), +# 一图出现多种缺陷时取框数最多的类别。 +# 无论取哪种,定位分支始终按 classes 逐类输出 heatmap。 +label_mode: multilabel + +image_size: [512, 512] +gray: false + +data: + train: + dataset: + name: film_defect + img_dir: data/film/images/train + label_dir: null # 缺省按 YOLO 约定取 data/film/labels/train + gray: false + max_objects: 64 + transforms: + - {name: Resize, size: [512, 512]} + - {name: RandomHorizontalFlip, p: 0.5} + - {name: RandomVerticalFlip, p: 0.5} + - {name: ColorJitter, brightness: 0.2, contrast: 0.2} + - {name: ToTensor} + - {name: Normalize, mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225]} + loader: + batch_size: 16 + shuffle: true + num_workers: 8 + drop_last: true + + val: + dataset: + name: film_defect + img_dir: data/film/images/val + label_dir: null + gray: false + transforms: + - {name: Resize, size: [512, 512]} + - {name: ToTensor} + - {name: Normalize, mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225]} + loader: + batch_size: 16 + shuffle: false + num_workers: 4 diff --git a/configs/eval/baseline.yaml b/configs/eval/baseline.yaml new file mode 100644 index 0000000..a1d385a --- /dev/null +++ b/configs/eval/baseline.yaml @@ -0,0 +1,12 @@ +_base_: + - ../train/baseline.yaml + +# 评估允许在 CPU 上跑(部署前的精度回归),训练则必须用 GPU。 +device: cuda +checkpoint: runs/baseline/best.pth + +eval: + split: val + score_threshold: 0.3 + iou_threshold: 0.5 + topk: 50 diff --git a/configs/model/resnet18.yaml b/configs/model/resnet18.yaml new file mode 100644 index 0000000..ce67959 --- /dev/null +++ b/configs/model/resnet18.yaml @@ -0,0 +1,15 @@ +# torchvision resnet18:stride 32 + 3 次上采样 => 输出 stride 4。 +model: + name: defect_net + backbone: + name: torchvision + arch: resnet18 + pretrained: true + in_channels: 3 + neck: + name: upsample + channels: 64 + num_upsamples: 3 + head: + hidden_channels: 64 + with_classification: true diff --git a/configs/model/tiny_conv.yaml b/configs/model/tiny_conv.yaml new file mode 100644 index 0000000..210c7ae --- /dev/null +++ b/configs/model/tiny_conv.yaml @@ -0,0 +1,18 @@ +# 自研轻量骨架:backbone stride 16 + 2 次上采样 => 输出 stride 4。 +model: + name: defect_net + backbone: + name: tiny_conv + in_channels: 3 + widths: [32, 64, 128, 256] + depths: [1, 2, 2, 2] + block: dw # dw=深度可分离(更快) / res=残差(更强) + neck: + name: upsample + channels: 64 + num_upsamples: 2 + head: + hidden_channels: 64 + with_classification: true + classification: + dropout: 0.1 diff --git a/configs/train/baseline.yaml b/configs/train/baseline.yaml new file mode 100644 index 0000000..a45bc4b --- /dev/null +++ b/configs/train/baseline.yaml @@ -0,0 +1,42 @@ +_base_: + - ../model/tiny_conv.yaml + - ../dataset/film_defect.yaml + +work_dir: runs/baseline +seed: 42 +deterministic: false +device: cuda # 训练只允许 Nvidia GPU,见 utils/device.py + +train: + epochs: 100 + amp: true + clip_grad: 5.0 + log_interval: 20 + monitor: f1 # 用于挑选 best.pth 的验证指标 + + optimizer: + name: adamw + lr: 0.001 + weight_decay: 0.0001 + + scheduler: + name: cosine + + loss: + heatmap_weight: 1.0 + wh_weight: 0.1 + offset_weight: 1.0 + cls_weight: 1.0 + + eval: + score_threshold: 0.3 + iou_threshold: 0.5 + topk: 50 + + # 训练过程可视化(SwanLab)。未安装或 enabled=false 时自动跳过。 + swanlab: + enabled: true + project: film-defect + name: null # 缺省用 work_dir 的目录名 + mode: online # online(云端) | local(本地看板) | offline(仅存盘) | disabled + tags: [tiny_conv, center-head] diff --git a/configs/train/multiclass.yaml b/configs/train/multiclass.yaml new file mode 100644 index 0000000..2bb1e3c --- /dev/null +++ b/configs/train/multiclass.yaml @@ -0,0 +1,12 @@ +_base_: + - ./baseline.yaml + +# 单类(互斥)分类示例:一张图只归一个类别,分类头多一维表示"无缺陷"。 +label_mode: multiclass + +work_dir: runs/multiclass +train: + monitor: cls_f1 # 单类模式下常用分类指标挑选 best.pth + swanlab: + name: multiclass + tags: [tiny_conv, multiclass] diff --git a/configs/train/resnet18.yaml b/configs/train/resnet18.yaml new file mode 100644 index 0000000..1d79d02 --- /dev/null +++ b/configs/train/resnet18.yaml @@ -0,0 +1,8 @@ +_base_: + - ./baseline.yaml + - ../model/resnet18.yaml + +work_dir: runs/resnet18 +train: + optimizer: + lr: 0.0005 diff --git a/eval.py b/eval.py new file mode 100755 index 0000000..e88b84e --- /dev/null +++ b/eval.py @@ -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()) diff --git a/train.py b/train.py new file mode 100755 index 0000000..4575945 --- /dev/null +++ b/train.py @@ -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())