Shaun ZhouandClaude Opus 5
4fc0fa5353
- docs/unsupervised-clustering.md: 逐图标准化对仿射光照群的不变性证明、 多 stage 特征的取舍、PCA 白化失效的谱分析、KMeans 收敛性与 kmeans++ 近似比、 轮廓系数与 NMI 稳定性的定义、NT-Xent 与 InfoNCE 下界，附全部消融数据与代码对照表 - figures/: 四张 mermaid 图（流水线、KMeans 控制流、SimCLR、主成分谱）， 含 .mmd 源、.md 预览与渲染好的 .png - README: 数据准备、训练评估、多类/单类、SwanLab、Profiler、聚类与自监督的运行说明 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
245 lines
12 KiB
Markdown
245 lines
12 KiB
Markdown
# ModelTrainEval
|
||
|
||
薄膜缺陷**分类 + 定位**模型的训练与评估。模型最终部署在实时设备上，因此只使用基于
|
||
convolution 的结构（不含 attention / transformer），并支持导出 ONNX。
|
||
|
||
## 目录
|
||
|
||
| 目录 | 作用 |
|
||
| --- | --- |
|
||
| `configs/` | yaml 配置（`model/` 模型、`dataset/` 数据、`train/` 训练、`eval/` 评估） |
|
||
| `models/` | 模型骨架：`backbones/` + `necks/` + `heads/`，由 `defect_net.py` 组装 |
|
||
| `datasets/` | 抛出数据样本的 Dataset 类、数据增强、监督信号编码 |
|
||
| `data/` | 数据集本体或软链接（不入库，见 [data/README.md](data/README.md)） |
|
||
| `scripts/` | 格式转换、数据集切分、ONNX 导出等批处理脚本 |
|
||
| `engine/` | 训练 / 评估循环、损失、指标、优化器、检查点、SwanLab 可视化 |
|
||
| `utils/` | 配置、注册表、日志、随机种子、GPU 守卫 |
|
||
| `tests/` | 单元测试 |
|
||
| `docs/` | 算法与数学说明文档 |
|
||
| `figures/` | 文档用的 mermaid 图（`.mmd` 源 + `.md` 预览 + `.png`） |
|
||
|
||
入口在仓库根目录：[train.py](train.py)、[eval.py](eval.py)。
|
||
|
||
## 环境
|
||
|
||
```bash
|
||
uv sync # 安装依赖（含 swanlab）
|
||
uv sync --group export # 需要导出 ONNX 时再装
|
||
uv run python scripts/check_env.py # 自检（Python / PyTorch / Nvidia GPU）
|
||
```
|
||
|
||
> **训练机必须有 Nvidia GPU。** 所有训练入口先调用 `utils.device.require_cuda()`，
|
||
> 检测不到 CUDA 构建的 PyTorch 或可用 GPU 时直接报错退出。评估与 ONNX 导出允许 CPU。
|
||
> Linux 上 PyPI 的 torch 即 CUDA 构建；如需指定 CUDA 版本：
|
||
> `uv pip install --index-url https://download.pytorch.org/whl/cu124 torch torchvision`
|
||
|
||
## 数据准备
|
||
|
||
标注为 **YOLO 风格**：`images/` 下每张图在同层 `labels/` 有一个同名 `.txt`，
|
||
每行 `class_id cx cy w h`（相对宽高归一化到 [0,1]），空文件或缺失 `.txt` 即无缺陷样本。
|
||
|
||
```
|
||
data/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/film/labelme --images data/film/images \
|
||
--output data/film/labels --data-yaml data/film/data.yaml
|
||
|
||
# 2. 切分 train / val（默认建软链，不复制数据）
|
||
uv run python scripts/split_dataset.py --images data/film/images \
|
||
--labels data/film/labels --val-ratio 0.2 --data-yaml data/film/data.yaml
|
||
```
|
||
|
||
配置里的 `classes` 既可以直接写类别列表，也可以写成 `data/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`（默认 `runs/<name>/`）：`train.log`、`config.yaml`、`last.pth`、`best.pth`、`swanlog/`。
|
||
|
||
## 多类 / 单类分类
|
||
|
||
图像级分类的形式由 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 记录
|
||
验证指标与显存峰值，实验的超参就是完整的 yaml 配置。在 `train.swanlab` 里配置：
|
||
|
||
```yaml
|
||
train:
|
||
swanlab:
|
||
enabled: true
|
||
project: film-defect
|
||
name: null # 缺省用 work_dir 目录名
|
||
mode: online # online(云端) | local(本地看板) | offline(仅存盘) | disabled
|
||
tags: [tiny_conv, center-head]
|
||
```
|
||
|
||
```bash
|
||
swanlab login # 首次使用云端模式需登录
|
||
uv run python train.py --config configs/train/baseline.yaml --set train.swanlab.mode=offline
|
||
uv run swanlab watch 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、CUDA 显存峰值（`max_memory_allocated` 与
|
||
`max_memory_reserved`）、权重占用，以及按显存和耗时各排一次序的算子级明细。
|
||
`--amp` 用 fp16 推理对比，`--postprocess` 把 heatmap 解码也计入。没有 GPU 时自动退回 CPU
|
||
（只有耗时与 CPU 内存）。
|
||
|
||
## 未分类样本的无监督聚类
|
||
|
||
`data/film` 里大量样本还没人工分完。这条流水线把目录下**所有图片拉平**（不看任何
|
||
目录名），做无监督聚类，产出可人工核对的分组：
|
||
|
||
```bash
|
||
# ImageNet 预训练卷积特征，不训练，CPU 上 14.5k 张约 100 秒
|
||
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml
|
||
|
||
# 先扫一遍簇数再决定分几类
|
||
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml --sweep-k 2-8
|
||
|
||
# 换成在本地数据上自监督训练的特征（SimCLR，属于训练，需要 Nvidia GPU）
|
||
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml --set method=simclr
|
||
```
|
||
|
||
簇数由配置的 `num_clusters` 决定，也可以 `--num-clusters 6` 临时覆盖（实测 4 类的换种子稳定性最高，见文档 §10.4）。产出在 `work_dir`：
|
||
`assignments.csv`（每张图的簇号与到质心的距离）、`summary.json`、
|
||
`clusters/cluster_XX/`（按簇建的软链目录，直接用看图软件翻）、
|
||
`cluster_XX_preview.png`（每簇最靠近质心的代表样本拼图）。
|
||
|
||
翻图时注意 `clusters/` 里文件名前缀是排序号：`0000_` 最靠近质心（最典型），
|
||
末尾 20 张是离质心最远的离群样本——最能暴露这个簇是不是混了东西。
|
||
|
||
**尺寸不统一的处理**：这批图是已裁好的缺陷小图，中位 152×113，宽高比从 0.24 到 39.9。
|
||
尺寸和宽高比是上游 ROI 裁剪的产物、不是缺陷属性，所以预处理的目标是把它抹掉：
|
||
直接拉伸到 128×128（不用 letterbox——填充比例等价于把宽高比明着喂给网络），
|
||
再做逐图 z-score 抹掉整体亮度/对比度。
|
||
|
||
无标签时用两个指标判断质量：**轮廓系数**（簇是否紧致）和**换随机种子的稳定性**
|
||
（同一份数据换初始化后划分是否还一样，低说明是在硬切一团连续数据）。
|
||
|
||
> 完整的数学推导、算法流程图与全部消融数据见
|
||
> **[docs/unsupervised-clustering.md](docs/unsupervised-clustering.md)**。
|
||
|
||
### 自监督预训练（SimCLR）
|
||
|
||
`method: pretrained`（默认）用 ImageNet 权重直接抽特征，不训练；`method: simclr` 则在
|
||
**本批薄膜图上自监督训练**一个项目自己的纯卷积骨架，域适配更好。同一张图的两个随机增强
|
||
视图互为正样本，batch 内其余 2B−2 个视图为负样本，用 NT-Xent 损失训练。
|
||
|
||
**这属于训练，因此必须有 Nvidia GPU**（走 `utils.device.require_cuda()` 守卫，
|
||
在无卡机器上会直接报错退出）：
|
||
|
||
```bash
|
||
# 1. 自监督训练 + 聚类一步跑完（首次运行）
|
||
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml --set method=simclr
|
||
|
||
# 2. 只想改簇数、复用已经训好的骨架：把权重路径填进 simclr.checkpoint
|
||
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml --set method=simclr simclr.checkpoint=runs/cluster/simclr_backbone.pth --num-clusters 6
|
||
|
||
# 3. 调训练超参
|
||
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml --set method=simclr simclr.epochs=200 simclr.batch_size=512 simclr.temperature=0.1
|
||
```
|
||
|
||
训练完权重存到 `<work_dir>/simclr_backbone.pth`（**只含骨架，不含投影头**——投影头被训练成
|
||
对增强不变，会丢掉下游有用的信息，所以聚类和保存都只用骨架特征）。相关配置：
|
||
|
||
```yaml
|
||
method: simclr
|
||
|
||
simclr:
|
||
# 自监督用项目自己的纯卷积骨架（与 method=pretrained 的顶层 `backbone` 各自独立，
|
||
# 两者字段完全不同，分开配置以免 _base_ 合并时串味）
|
||
backbone: {name: tiny_conv, in_channels: 3, widths: [32, 64, 128, 256], depths: [1, 2, 2, 2]}
|
||
epochs: 100
|
||
batch_size: 256
|
||
temperature: 0.2 # 越小越关注难负样本
|
||
amp: true
|
||
checkpoint: null # 填已有权重可跳过重复训练
|
||
projection: {hidden_channels: 256, out_channels: 128}
|
||
optimizer: {name: adamw, lr: 0.001, weight_decay: 0.0001}
|
||
scheduler: {name: cosine}
|
||
|
||
view_transforms: # 生成两个视图的随机增强
|
||
- {name: RandomResizedCrop, size: [128, 128], scale: [0.5, 1.0]}
|
||
- {name: RandomHorizontalFlip, p: 0.5}
|
||
- {name: RandomVerticalFlip, p: 0.5}
|
||
- {name: GaussianBlur, p: 0.3}
|
||
```
|
||
|
||
> `RandomResizedCrop` 的 `scale` 下限别低于 0.4：裁太狠会让两个视图都只剩背景，
|
||
> "背景 vs 背景"被当成正样本对，学出来的特征会退化。
|
||
|
||
训练出的骨架和检测模型用的是同一套卷积骨架，可以直接拿去初始化 `DefectNet` 的 backbone——
|
||
标注量有限时，通常比 ImageNet 权重更贴合薄膜的成像特性。
|
||
|
||
## 方法
|
||
|
||
- **骨架**：`tiny_conv`（自研深度可分离卷积，默认）或 `torchvision`（resnet / mobilenet / shufflenet）。
|
||
- **neck**：nearest 上采样 + 3×3 卷积，把特征恢复到 stride 4，兼顾小缺陷定位与速度。
|
||
- **定位头**：anchor-free 中心点（CenterNet 风格）——`heatmap` + `wh` + `offset`，
|
||
推理端用 3×3 max-pool 做峰值抑制，无 NMS 循环，便于导出到 TensorRT / OpenVINO。
|
||
- **分类头**：全局池化 + 全连接，输出**多热**标签（一张图可同时含多种缺陷，全 0 表示 OK 片）。
|
||
- **指标**：定位 precision / recall / F1 / mean IoU（IoU 阈值匹配），分类 micro-F1 与整图完全正确率。
|
||
`train.monitor` 决定用哪个指标挑选 `best.pth`。
|
||
|
||
配置支持 `_base_` 继承与命令行 `--set a.b.c=value` 覆盖，见 [utils/config.py](utils/config.py)。
|
||
新增骨架 / 数据集只需用注册表装饰器注册，再在 yaml 里按 `name` 引用。
|
||
|
||
## 开发
|
||
|
||
```bash
|
||
uv run pytest # 单元测试
|
||
uv run ruff check . # 静态检查
|
||
```
|