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:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
ef64e5bf77
commit
b0d0dfffbb
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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]
|
||||
@@ -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]
|
||||
@@ -0,0 +1,8 @@
|
||||
_base_:
|
||||
- ./baseline.yaml
|
||||
- ../model/resnet18.yaml
|
||||
|
||||
work_dir: runs/resnet18
|
||||
train:
|
||||
optimizer:
|
||||
lr: 0.0005
|
||||
@@ -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())
|
||||
@@ -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())
|
||||
Reference in New Issue
Block a user