refactor: 按实际的五步工作流裁掉无关代码
实际工作流是：simclr 训练 -> 聚类 -> 人工归类 -> 分类训练 -> GradCAM/SHAP 生成
热力图与 box。据此移除不参与其中的检测/定位链路，共 31 个文件、约 2600 行：
- 模型：defect_net、postprocess、necks/、heads/center_head
- 数据：film_defect（YOLO 读取端）、target_encoder
- 引擎：trainer、evaluator、losses
- 入口：train.py、eval.py
- 脚本：convert_annotations、split_dataset、export_onnx、profile_inference
- 配置：configs/{train,eval,model,dataset}
- utils/labels（多类/单类标签模式，仅检测分类头使用）
共享模块按实际用量瘦身：engine/metrics 只留 AverageMeter（检测指标随之移除）；
datasets/yolo 只留写出侧（IMAGE_SUFFIXES 与 write_data_yaml）——第 5 步产出的
YOLO 标签是给外部检测训练用的交换格式，读取端不再需要。
裁剪过程中修掉一个真 bug：generate_boxes 的 --stage 默认 6（layer3）对 tiny_conv
这类非 nn.Sequential 骨架直接抛异常。现在自动退回骨架最终输出并警告，
另加 --stage -1 显式选择最终输出。
五步工作流已端到端实跑验证（simclr 1 轮 -> 聚类 -> 导出 -> 分类 1 轮 ->
gradcam 与 shap 两种方法各生成一遍）。
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
57446c2f10
commit
86c86b105c
@@ -1,7 +1,15 @@
|
||||
# ModelTrainEval
|
||||
|
||||
薄膜缺陷**分类 + 定位**模型的训练与评估。模型最终部署在实时设备上，因此只使用基于
|
||||
convolution 的结构（不含 attention / transformer），并支持导出 ONNX。
|
||||
薄膜缺陷的**无监督归类 → 分类模型 → 弱监督定位**流水线。模型最终部署在实时设备上，
|
||||
因此只使用基于 convolution 的结构（不含 attention / transformer）。
|
||||
|
||||
完整工作流五步：
|
||||
|
||||
1. **自监督**：在无标注数据上用 SimCLR 训练卷积骨架
|
||||
2. **聚类**：按指定簇数分组，产出可人工翻看的硬链接目录
|
||||
3. **人工归类**：在 `runs/` 下直接整理簇（可对某一类继续聚类细化），再导出成数据集
|
||||
4. **分类训练**：在归类好的数据集上训练分类模型
|
||||
5. **定位**：用 Grad-CAM / SHAP 生成热力图与 YOLO 框标签
|
||||
|
||||
## 目录
|
||||
|
||||
@@ -18,7 +26,7 @@ convolution 的结构（不含 attention / transformer），并支持导出 ONNX
|
||||
| `docs/` | 算法与数学说明文档 |
|
||||
| `figures/` | 文档用的 mermaid 图（`.mmd` 源 + `.md` 预览 + `.png`） |
|
||||
|
||||
入口在仓库根目录：[train.py](train.py)（检测+定位训练）、[eval.py](eval.py)、[classify.py](classify.py)（分类训练与评估）。
|
||||
入口：[classify.py](classify.py)（分类训练与评估）、[scripts/](scripts/) 下的聚类、导出与框生成脚本。
|
||||
|
||||
## 环境
|
||||
|
||||
@@ -50,98 +58,6 @@ uv run python train.py --config configs/train/baseline.yaml --device cpu
|
||||
`amp` 字段接受 `auto`（随后端）/ `true` / `false`。MPS 上 `torch.profiler` 抓不到 Metal kernel
|
||||
（没有对应的 `ProfilerActivity`），算子明细只反映 CPU 侧下发耗时。
|
||||
|
||||
## 数据准备
|
||||
|
||||
标注为 **YOLO 风格**：`images/` 下每张图在同层 `labels/` 有一个同名 `.txt`，
|
||||
每行 `class_id cx cy w h`（相对宽高归一化到 [0,1]），空文件或缺失 `.txt` 即无缺陷样本。
|
||||
|
||||
```
|
||||
data/raw/film/
|
||||
├── images/{train,val}/0001.png
|
||||
├── labels/{train,val}/0001.txt
|
||||
└── data.yaml # path / train / val / nc / names
|
||||
```
|
||||
|
||||
```bash
|
||||
# 1. labelme / VOC -> YOLO txt（顺带生成 data.yaml）
|
||||
uv run python scripts/convert_annotations.py --format labelme \
|
||||
--input data/raw/film/labelme --images data/raw/film/images \
|
||||
--output data/raw/film/labels --data-yaml data/raw/film/data.yaml
|
||||
|
||||
# 2. 切分 train / val（默认建软链，不复制数据）
|
||||
uv run python scripts/split_dataset.py --images data/raw/film/images \
|
||||
--labels data/raw/film/labels --val-ratio 0.2 --data-yaml data/raw/film/data.yaml
|
||||
```
|
||||
|
||||
配置里的 `classes` 既可以直接写类别列表，也可以写成 `data/raw/film/data.yaml`，由框架读 `names`。
|
||||
读取时也接受 YOLO 分割格式的多边形行（取外接矩形）。细节见 [data/README.md](data/README.md)
|
||||
与 [datasets/yolo.py](datasets/yolo.py)。
|
||||
|
||||
## 训练与评估
|
||||
|
||||
```bash
|
||||
uv run python train.py --config configs/train/baseline.yaml
|
||||
uv run python train.py --config configs/train/baseline.yaml --set train.epochs=10 train.optimizer.lr=5e-4
|
||||
uv run python train.py --config configs/train/baseline.yaml --resume runs/baseline/last.pth
|
||||
|
||||
uv run python eval.py --config configs/eval/baseline.yaml --checkpoint runs/baseline/best.pth
|
||||
uv run python scripts/export_onnx.py --config configs/train/baseline.yaml \
|
||||
--checkpoint runs/baseline/best.pth --output runs/baseline/model.onnx
|
||||
```
|
||||
|
||||
换骨架用对应的配置对（训练与评估必须成对，否则构建出的骨架与权重对不上）：
|
||||
|
||||
| 骨架 | 训练 | 评估 | work_dir |
|
||||
| --- | --- | --- | --- |
|
||||
| `tiny_conv`（默认） | `configs/train/baseline.yaml` | `configs/eval/baseline.yaml` | `runs/baseline` |
|
||||
| `resnet18` | `configs/train/resnet18.yaml` | `configs/eval/resnet18.yaml` | `runs/resnet18` |
|
||||
|
||||
```bash
|
||||
uv run python train.py --config configs/train/resnet18.yaml
|
||||
uv run python eval.py --config configs/eval/resnet18.yaml
|
||||
uv run python scripts/export_onnx.py --config configs/train/resnet18.yaml \
|
||||
--checkpoint runs/resnet18/best.pth --output runs/resnet18/model.onnx
|
||||
```
|
||||
|
||||
要加别的骨架（mobilenet / shufflenet），照 [configs/model/resnet18.yaml](configs/model/resnet18.yaml)
|
||||
复制一份改 `arch` 即可——注意保留里面的 `_replace_: true`，否则 `tiny_conv` 的
|
||||
`widths/depths/block` 会残留下来传给新骨架。
|
||||
|
||||
产物写在 `work_dir`（默认 `runs/<name>/`）：`train.log`、`config.yaml`、`last.pth`、`best.pth`、`swanlog/`。
|
||||
|
||||
**续训以配置为准。** `--resume` 只从检查点取"进度"——模型权重、优化器动量缓冲、
|
||||
已训到第几轮、`best_metric`、`global_step`；`epochs`、`lr`、`weight_decay`、调度器
|
||||
一律以当前 yaml 为准，所以可以顺手调超参：
|
||||
|
||||
```bash
|
||||
uv run python train.py --config configs/train/baseline.yaml --resume runs/baseline/last.pth \
|
||||
--set train.epochs=200 train.optimizer.lr=5e-4
|
||||
```
|
||||
|
||||
与检查点不一致时日志会逐项列出（`配置与检查点不一致，按当前配置执行：总轮数 100 -> 200；
|
||||
optimizer.lr 0.001 -> 0.0005`）。学习率调度按新的总轮数重建并快进到当前 epoch。
|
||||
`epochs` 是**总轮数**不是剩余轮数。SimCLR 自监督那边是同样的语义。
|
||||
|
||||
## 多类 / 单类分类
|
||||
|
||||
图像级分类的形式由 yaml 的 `label_mode` 决定，模型的分类头、分类损失、分类指标都跟着切换
|
||||
（定位分支不受影响，始终按 `classes` 逐类输出 heatmap）：
|
||||
|
||||
| `label_mode` | 含义 | 分类头输出 | 损失 | 主要指标 |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| `multilabel`（默认） | **多类**：一张图可同时含多种缺陷，全 0 表示 OK 片 | `len(classes)` | 逐类 BCE | micro P/R/F1、`cls_exact_match` |
|
||||
| `multiclass` | **单类**：一张图只归一个类别，互斥 | `len(classes) + 1`（末位为"无缺陷"） | softmax CE | `cls_accuracy`、macro P/R/F1 |
|
||||
|
||||
```bash
|
||||
uv run python train.py --config configs/train/baseline.yaml # 多类
|
||||
uv run python train.py --config configs/train/multiclass.yaml # 单类
|
||||
uv run python train.py --config configs/train/baseline.yaml --set label_mode=multiclass
|
||||
```
|
||||
|
||||
单类模式下若一张图标注了多种缺陷，取框数最多的类别（并列取下标最小的）；
|
||||
无缺陷样本落到背景下标 `len(classes)`。两种模式的分类指标都含 `cls_f1`，
|
||||
因此 `train.monitor` 不必跟着改。
|
||||
|
||||
## 训练可视化（SwanLab）
|
||||
|
||||
训练过程通过 [SwanLab](https://swanlab.cn) 记录：按 step 记录各项损失与学习率，按 epoch 记录
|
||||
@@ -176,26 +92,6 @@ uv run swanlab sync runs/baseline/swanlog
|
||||
未安装 swanlab、未登录或网络不通时会打一条警告并继续训练，不会中断。关掉用
|
||||
`--set train.swanlab.enabled=false`。
|
||||
|
||||
## 推理性能与显存分析
|
||||
|
||||
部署前用 `torch.profiler` 看单帧耗时与显存占用：
|
||||
|
||||
```bash
|
||||
uv run python scripts/profile_inference.py --config configs/train/baseline.yaml \
|
||||
--checkpoint runs/baseline/best.pth --batch-size 1 4 8 --postprocess
|
||||
|
||||
# 导出 chrome trace，用 chrome://tracing 或 ui.perfetto.dev 打开
|
||||
uv run python scripts/profile_inference.py --config configs/train/baseline.yaml \
|
||||
--checkpoint runs/baseline/best.pth --trace runs/baseline/trace.json --out runs/baseline/profile.json
|
||||
```
|
||||
|
||||
输出包含：每个 batch 的单帧延迟 / FPS、显存峰值、权重占用，以及按显存和耗时各排一次序的
|
||||
算子级明细。`--amp` 用 fp16 推理对比，`--postprocess` 把 heatmap 解码也计入，
|
||||
`--device` 指定后端（默认 auto）。
|
||||
|
||||
后端差异见[计算后端](#计算后端)：CUDA 给精确峰值；MPS 上 `torch.profiler` 抓不到 Metal kernel，
|
||||
算子明细只反映 CPU 侧下发耗时、显存靠采样是下界；CPU 只有耗时与 CPU 内存。
|
||||
|
||||
## 未分类样本的无监督聚类
|
||||
|
||||
### 调用 ImageNet 预训练模型
|
||||
|
||||
Reference in New Issue
Block a user