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
2026-09-10 14:58:41 +08:00
co-authored by Claude Opus 5
parent 33d28dcc7c
commit 33fdf396ee
10 changed files with 1085 additions and 1 deletions
+48 -1
View File
@@ -18,7 +18,7 @@ convolution 的结构(不含 attention / transformer),并支持导出 ONNX
| `docs/` | 算法与数学说明文档 |
| `figures/` | 文档用的 mermaid 图(`.mmd` 源 + `.md` 预览 + `.png`) |
入口在仓库根目录:[train.py](train.py)、[eval.py](eval.py)。
入口在仓库根目录:[train.py](train.py)(检测+定位训练)、[eval.py](eval.py)、[classify.py](classify.py)(分类训练与评估)。
## 环境
@@ -585,6 +585,53 @@ uv run swanlab watch runs/simclr/swanlog
训练出的骨架和检测模型用的是同一套卷积骨架,可以直接拿去初始化 `DefectNet` 的 backbone——
标注量有限时,通常比 ImageNet 权重更贴合薄膜的成像特性。
## 缺陷分类模型
在聚类 + 人工归类得到的数据集上训练分类模型(`data/<数据集>/<类别名>/*.png`):
```bash
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
```
产物在 `work_dir`:`best.pth` / `last.pth` / `config.yaml` / `classify.log` /
`eval_report.json`(逐类 P/R/F1 与完整混淆矩阵)。
### 三个关键取舍
**主指标是 macro-F1,不是准确率。** 数据从 `normal` 2679 张到 `others` 65 张差 41 倍,
全预测成 normal 也能拿 24% 准确率。配置里同时给出 `balanced_accuracy`(各类召回的均值),
两者一起看才知道小类有没有学会。
**损失按类别频率加权**(`class_weight: balanced`,逆频率归一化到均值 1),
否则占 24% 的 normal 会主导梯度。不需要时设 `none`。
**逐图标准化抹掉整体亮度。** 15 个类别里有几类几乎只来自单一原始数据集
(`ripple` 几乎全是 pearlescent,`grease`/`zigzag` 全是 film),不抹掉亮度的话模型
可能靠"认出这是哪个数据集"就把类别猜对——验证集上分数很高,换批数据就崩。
### 划分
没有预先的 train/val,按**类别分层**切分(默认 8:2,固定种子)。逐类切分是必须的:
`others` 只有 65 张,整体随机切很容易让它在验证集里一张不剩。
### 用自监督骨架初始化
前面自监督训练出的骨架可以直接拿来初始化——这正是聚类阶段产物的下游用途:
```yaml
model:
backbone: {name: torchvision, arch: resnet101} # 必须与自监督时一致
init_from: runs/simclr_r101/simclr_backbone.pth
frozen_backbone: false # true = 冻结骨架只训分类头(线性探针)
```
只加载 backbone,分类头仍随机初始化。`frozen_backbone: true` 可以单独衡量
自监督特征本身有多好。
## 方法
- **骨架**:`tiny_conv`(自研深度可分离卷积,默认)或 `torchvision`(resnet / mobilenet / shufflenet)。