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
2026-09-02 17:55:20 +08:00
co-authored by Claude Opus 5
parent ef64e5bf77
commit b0d0dfffbb
9 changed files with 344 additions and 0 deletions
+51
View File
@@ -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
+12
View File
@@ -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
+15
View File
@@ -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
+18
View File
@@ -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
+42
View File
@@ -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]
+12
View File
@@ -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]
+8
View File
@@ -0,0 +1,8 @@
_base_:
- ./baseline.yaml
- ../model/resnet18.yaml
work_dir: runs/resnet18
train:
optimizer:
lr: 0.0005