原始数据从 data/<数据集> 移到了 data/raw/<数据集>。
- 聚类配置的 dataset.root:film / pearlescent 均改为 data/raw/ 下
- 检测配置的 img_dir 与 label_dir 注释同步
- README、data/README、算法文档、流程图(含重新渲染的 PNG)与各脚本
docstring 里的路径示例一并更新
- data/README 补充目录约定:原始数据只读放 data/raw/,
export_clusters.py 的派生数据集放在 data/ 下与 raw/ 平级
- 顺手把 unlabeled 数据集加进 datasets/__init__.py:此前它没在包初始化时注册,
build_dataset({"name": "unlabeled"}) 只有在别处先 import 过该模块时才可用
迁移后重跑全量聚类,结果与迁移前逐项一致(轮廓 0.3255、稳定性 0.9641、
簇大小 [4507, 3795, 4849, 1428])。
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
534 lines
26 KiB
Markdown
534 lines
26 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/` | 配置、注册表、日志、随机种子、设备解析(cuda/mps/cpu) |
|
||
| `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 / 可用后端)
|
||
```
|
||
|
||
### 计算后端
|
||
|
||
支持 **CUDA / MPS(Apple Silicon)/ CPU**。配置里的 `device: auto` 按 `cuda > mps > cpu`
|
||
自动挑最快的可用后端,也可以显式写 `cuda` / `mps` / `cpu`,或用命令行 `--device` 覆盖:
|
||
|
||
```bash
|
||
uv run python train.py --config configs/train/baseline.yaml # auto
|
||
uv run python train.py --config configs/train/baseline.yaml --device mps # 强制 Mac GPU
|
||
uv run python train.py --config configs/train/baseline.yaml --device cpu
|
||
```
|
||
|
||
**显式指定不可用的后端会直接报错退出**,不会悄悄退回 CPU——训练慢几十倍却不吭声是最坏的情况。
|
||
|
||
| 后端 | AMP 默认 | 显存统计 | 备注 |
|
||
| --- | --- | --- | --- |
|
||
| `cuda` | **开** | 精确峰值 | Linux 上 PyPI 的 torch 即 CUDA 构建;指定版本用 `uv pip install --index-url https://download.pytorch.org/whl/cu124 torch torchvision` |
|
||
| `mps` | **关** | 采样下界 | Apple Silicon。autocast 可用且数值正常,但实测比 fp32 慢(tiny_conv @128²:fp32 76 ms/步,fp16 90 ms),故默认不开;`amp: true` 可强制打开 |
|
||
| `cpu` | 不支持 | 无 | 实测比 MPS 慢约 30 倍,仅适合调试与精度回归 |
|
||
|
||
`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/`。
|
||
|
||
## 多类 / 单类分类
|
||
|
||
图像级分类的形式由 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
|
||
# 本地看板:不需要账号,训练时用 local 模式,另开终端起看板
|
||
uv run python train.py --config configs/train/baseline.yaml --set train.swanlab.mode=local
|
||
uv run swanlab watch runs/baseline/swanlog # 默认 http://127.0.0.1:5092
|
||
|
||
# 云端
|
||
uv run swanlab login # 首次使用需登录
|
||
uv run python train.py --config configs/train/baseline.yaml # mode 默认 online
|
||
|
||
# 离线跑完后再补传云端
|
||
uv run swanlab sync runs/baseline/swanlog
|
||
```
|
||
|
||
看板可以在训练**进行中**打开,指标边训边刷。端口冲突用 `-p 5100`,要局域网访问加 `-h 0.0.0.0`。
|
||
本地看板依赖 `swanlab[dashboard]`,已包含在 `uv sync` 里。
|
||
|
||
未安装 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 预训练模型
|
||
|
||
`data/raw/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,Mac 上自动走 MPS)
|
||
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`(每簇最靠近质心的代表样本拼图)、
|
||
`cluster_scatter.png`(**簇分布图**,见下)。
|
||
|
||
### 簇分布图
|
||
|
||
`cluster_scatter.png` 和 preview 拼图同级,一张图回答"各簇离多远":
|
||
|
||
- **左**:聚类空间的 PCA 二维投影,散点按簇着色,`X` 是质心。标题给出两个主成分的
|
||
方差占比——占比低说明高维结构被压掉得多,图上的远近只能参考。
|
||
- **右**:质心两两距离矩阵(在**聚类空间**里算的,不是投影后),带数值标注。
|
||
这个是准的,看重叠严不严重以它为准。
|
||
|
||
投影用 PCA 而不是 t-SNE/UMAP:后两者为了把簇拉开会**扭曲簇间距离**,而这张图的目的
|
||
恰恰是看距离。相关配置:
|
||
|
||
```yaml
|
||
scatter:
|
||
enabled: true
|
||
max_points: 5000 # 点太多会糊成一团且文件很大,超过就随机抽样
|
||
point_size: 6
|
||
alpha: 0.35
|
||
figsize: [14, 6]
|
||
dpi: 150
|
||
```
|
||
|
||
关掉用 `--set scatter.enabled=false`。图上的中文需要系统里有 PingFang / Noto Sans CJK
|
||
等字体,找不到会自动退回英文标签(不会画成豆腐块)。质心距离矩阵同时写进
|
||
`summary.json` 的 `centroid_distances`,方便脚本读取。
|
||
|
||
翻图时注意 `clusters/` 里文件名前缀是排序号:`0000_` 最靠近质心(最典型),
|
||
末尾 20 张是离质心最远的离群样本——最能暴露这个簇是不是混了东西。
|
||
|
||
**尺寸不统一的处理**:这批图是已裁好的缺陷小图,中位 152×113,宽高比从 0.24 到 39.9。
|
||
尺寸和宽高比是上游 ROI 裁剪的产物、不是缺陷属性,所以预处理的目标是把它抹掉:
|
||
直接拉伸到 128×128(不用 letterbox——填充比例等价于把宽高比明着喂给网络),
|
||
再做逐图 z-score 抹掉整体亮度/对比度。
|
||
|
||
无标签时用两个指标判断质量:**轮廓系数**(簇是否紧致)和**换随机种子的稳定性**
|
||
(同一份数据换初始化后划分是否还一样,低说明是在硬切一团连续数据)。
|
||
|
||
> 完整的数学推导、算法流程图与全部消融数据见
|
||
> **[docs/unsupervised-clustering.md](docs/unsupervised-clustering.md)**。
|
||
|
||
### 按聚类结果导出数据集
|
||
|
||
人工核对完分组后,把结果落成一个可以直接使用的数据集目录:
|
||
|
||
```bash
|
||
uv run python scripts/export_clusters.py --work-dir runs/simclr_r18 # 先看计划加 --dry-run
|
||
```
|
||
|
||
```
|
||
data/film_r18_cluster_k4/ # <源目录名>_<骨架标签>_cluster_k<簇数>,自动推导
|
||
├── 0/ # 簇 0 的图,文件名保持原样
|
||
├── 1/
|
||
├── 2/
|
||
├── 3/
|
||
└── manifest.json # 来源 run、特征配置、每张图的原始路径
|
||
```
|
||
|
||
**复制真实文件**,不是软链也不改名,所以这个目录可以直接拷走或喂给别的工具。
|
||
目录名可用 `--tag r18ss` 换标签、`--out` 完全指定。几个保护:源文件缺失、
|
||
同簇内出现重名(保持原名会互相覆盖)都会直接报错;目标目录非空时必须显式 `--overwrite`。
|
||
|
||
> 标签只反映骨架,不区分 pretrained 与 simclr——同一骨架的两种跑法会推出同名目录,
|
||
> 需要并存时用 `--tag` 手动区分。无论目录叫什么,`manifest.json` 都记着这批图来自哪次 run。
|
||
|
||
### 自监督预训练(SimCLR)
|
||
|
||
`method: pretrained`(默认)用 ImageNet 权重直接抽特征,不训练;`method: simclr` 则在
|
||
**本批薄膜图上自监督训练**一个项目自己的纯卷积骨架,域适配更好。同一张图的两个随机增强
|
||
视图互为正样本,batch 内其余 2B−2 个视图为负样本,用 NT-Xent 损失训练。
|
||
|
||
训练走 `device: auto` 选中的后端——**Mac 上直接用 MPS 即可**,实测 14,579 张、
|
||
batch 256 时约 65 秒/epoch。
|
||
|
||
#### 第一步:训练 + 聚类(一条命令)
|
||
|
||
`method=simclr` 会先训练、存权重,然后直接用这份权重抽特征做聚类,不需要分两次调用:
|
||
|
||
```bash
|
||
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
|
||
--work-dir runs/simclr \
|
||
--set method=simclr simclr.epochs=40 loader.num_workers=8
|
||
```
|
||
|
||
按 65 秒/epoch 算,40 轮约 **45 分钟**。产出:
|
||
|
||
```
|
||
runs/simclr/
|
||
├── simclr_backbone.pth # 只含骨架,给下游用(聚类 / 初始化 DefectNet)
|
||
├── simclr_last.pth # 完整训练状态,给断点续训用
|
||
├── swanlog/ # SwanLab 日志
|
||
├── cluster.log
|
||
├── assignments.csv
|
||
├── summary.json
|
||
├── clusters/cluster_XX/
|
||
└── cluster_XX_preview.png
|
||
```
|
||
|
||
跑长任务建议丢后台,随时能看进度:
|
||
|
||
```bash
|
||
nohup uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
|
||
--work-dir runs/simclr \
|
||
--set method=simclr simclr.epochs=40 loader.num_workers=8 > runs/simclr.log 2>&1 &
|
||
|
||
tail -f runs/simclr.log
|
||
```
|
||
|
||
#### 断点续训
|
||
|
||
训练中断(Ctrl-C、kill、机器重启)后接着练,用 `simclr.resume` 指向 `simclr_last.pth`:
|
||
|
||
```bash
|
||
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
|
||
--work-dir runs/simclr \
|
||
--set method=simclr simclr.epochs=40 loader.num_workers=8 \
|
||
simclr.resume=runs/simclr/simclr_last.pth
|
||
```
|
||
|
||
日志会打印"已载入训练状态 …:从 epoch N 接着练到 40"。恢复的不只是权重,还有**投影头、
|
||
优化器动量缓冲和 step 计数**——少了这些,续训等于把动量清零,损失会先反弹一段。
|
||
|
||
> **配置文件是唯一真相源。** 存档只提供"进度"(权重、动量、已训到第几轮);
|
||
> `epochs`、`lr`、`weight_decay`、调度器一律以当前 yaml 为准。所以续训时可以顺手调超参:
|
||
>
|
||
> ```bash
|
||
> uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
|
||
> --work-dir runs/simclr \
|
||
> --set method=simclr simclr.epochs=60 simclr.optimizer.lr=5e-4 \
|
||
> simclr.resume=runs/simclr/simclr_last.pth
|
||
> ```
|
||
>
|
||
> 与存档不一致时日志会逐项列出(`配置与存档不一致,按当前配置执行:总轮数 40 -> 60;
|
||
> optimizer.lr 0.001 -> 0.0005`),不会出现"改了没生效"。学习率调度按新的总轮数重建
|
||
> 并快进到当前 epoch,不会沿用存档里的 T_max。
|
||
|
||
续训时各项配置的行为:
|
||
|
||
| 配置项 | 续训时 |
|
||
| --- | --- |
|
||
| `epochs` | 按 yaml,且是**总轮数**不是剩余轮数。已训 12 轮、写 40 就再训 28 轮;写成 ≤12 则不再训练并给出提示 |
|
||
| `optimizer` 的超参(`lr` / `weight_decay` / `betas` / `momentum` …) | 按 yaml |
|
||
| `scheduler`(类型与参数) | 按 yaml,重建后快进到当前 epoch |
|
||
| `temperature`、`log_interval`、`save_interval`、`amp`、`batch_size`、增强 | 按 yaml |
|
||
| 模型权重、投影头、优化器动量缓冲、已训轮数、`global_step` | 来自存档(这些才是"进度") |
|
||
| `optimizer.name`(adamw ↔ sgd) | 可以改,但动量缓冲无法迁移、会从零积累(有日志提示) |
|
||
| `backbone` 结构、`projection` 维度 | **不能改**,与存档权重形状不符会直接报错 |
|
||
|
||
> `resume` 与 `checkpoint` 语义不同,二选一(同时给会报错):
|
||
> **`resume`** = 载入完整状态**接着训练**(要 `simclr_last.pth`);
|
||
> **`checkpoint`** = 有权重就**跳过训练**直接聚类(要 `simclr_backbone.pth`)。
|
||
> 拿 `simclr_backbone.pth` 去 `resume` 会报错,因为它不含投影头和优化器状态。
|
||
|
||
#### 第二步:复用权重重新聚类
|
||
|
||
想换簇数或改聚类参数时**不必重训**——把权重路径填进 `simclr.checkpoint` 就会跳过训练
|
||
(日志会打印"复用已有的自监督骨架"),这一步只跑特征提取 + 聚类,约 **1 分钟**:
|
||
|
||
```bash
|
||
# 换成 k=2(实测 SimCLR 特征在 k=2 上最干净:轮廓 0.469 / 稳定性 0.991)
|
||
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
|
||
--work-dir runs/simclr_k2 \
|
||
--set method=simclr simclr.checkpoint=runs/simclr/simclr_backbone.pth \
|
||
--num-clusters 2
|
||
|
||
# 先扫一遍簇数再决定(只打指标表,不产出分组)
|
||
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
|
||
--set method=simclr simclr.checkpoint=runs/simclr/simclr_backbone.pth \
|
||
--sweep-k 2-8
|
||
```
|
||
|
||
> 复用权重时给一个**新的 `--work-dir`**,否则新的聚类产出会覆盖同目录下的旧结果。
|
||
|
||
**实测效果**(14,579 张,与 pretrained 用完全相同的聚类设置):只训 **1 个 epoch** 的域内特征
|
||
在 k=4 上就追平 ImageNet(轮廓 0.334 vs 0.325),在 **k=2 上大幅胜出**(轮廓 0.469、
|
||
稳定性 0.991)——目视确认这个二分正是"有明显缺陷 / 弱纹理近空白",可直接当作流水线的
|
||
第一级过滤。详见 [文档 §10.5](docs/unsupervised-clustering.md#105-simclr-域内自监督-vs-imagenet-预训练)。
|
||
|
||
权重存到 `<work_dir>/simclr_backbone.pth`(**只含骨架,不含投影头**——投影头被训练成
|
||
对增强不变,会丢掉下游有用的信息,所以聚类和保存都只用骨架特征)。
|
||
|
||
#### 固化到配置
|
||
|
||
每次敲 `--set` 麻烦的话,直接改 [configs/cluster/film.yaml](configs/cluster/film.yaml):
|
||
|
||
```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: 40
|
||
batch_size: 256
|
||
temperature: 0.2 # 越小越关注难负样本
|
||
amp: auto # CUDA 上开、MPS 上关
|
||
save_interval: 5 # 每 5 个 epoch 落一次盘
|
||
# 下面两个二选一,语义不同:
|
||
checkpoint: null # 有权重就跳过训练直接聚类,例 runs/simclr/simclr_backbone.pth
|
||
resume: null # 载入完整状态接着训练,例 runs/simclr/simclr_last.pth
|
||
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 背景"被当成正样本对,学出来的特征会退化。
|
||
|
||
#### 换一个骨架
|
||
|
||
**仓库自带 resnet18 / resnet50 的现成配置**,直接用:
|
||
|
||
```bash
|
||
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film_simclr_resnet18.yaml \
|
||
--work-dir runs/simclr_r18 --set simclr.epochs=40
|
||
|
||
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film_simclr_resnet50.yaml \
|
||
--work-dir runs/simclr_r50 --set simclr.epochs=40
|
||
```
|
||
|
||
两者在 Apple MPS 上的实测开销(14,579 张,128×128 双视图):
|
||
|
||
| 骨架 | 参数量 | layer2+3 维度 | 每 epoch | batch 上限 |
|
||
| --- | --- | --- | --- | --- |
|
||
| resnet18 | 11.2M | 384 | 约 93 秒 | 1024 可用 |
|
||
| resnet50 | 23.5M | 1536 | 约 155 秒 | **1024 会 OOM**(需 18.7 GiB),配置里用 256 |
|
||
|
||
resnet50 每 epoch 的耗时与 batch 基本无关(算力受限),所以显存不够时往下调 batch 不花额外时间,
|
||
只是 NT-Xent 的负样本变少。
|
||
|
||
注意默认的 `method: pretrained` 用的**也是 resnet18**(ImageNet 权重,不训练,约 51 秒出结果);
|
||
上面这个配置是在 `data/raw/film` 上自监督训练它。两者实测对比(14,579 张,k=4):
|
||
|
||
| 方案 | 轮廓 | 稳定性 | 成本 |
|
||
| --- | --- | --- | --- |
|
||
| pretrained resnet18(ImageNet,不训练) | **0.326** | **0.964** | 约 51 秒 |
|
||
| SimCLR tiny_conv(1 epoch) | 0.334 | 0.873 | 65 秒/epoch |
|
||
| SimCLR resnet18(1 epoch) | 0.243 | 0.997 | 93 秒/epoch,显存 315 MB |
|
||
|
||
resnet18 参数量是 tiny_conv 的 37 倍(11.2M vs 0.3M),1 个 epoch 远不够收敛——
|
||
上表只是**下界**,要用它就得训几十轮。想快速出结果,默认的 pretrained 或 tiny_conv 更划算。
|
||
|
||
`simclr.backbone` 用的是和检测模型同一套注册表,可选:
|
||
|
||
| `name` | `arch` | out_channels | 说明 |
|
||
| --- | --- | --- | --- |
|
||
| `tiny_conv` | — | 256(可调) | 自研深度可分离卷积,默认,最轻 |
|
||
| `torchvision` | `resnet18` / `resnet34` | 512 | 通用基线 |
|
||
| `torchvision` | `resnet50` | 2048 | 最强也最重,MPS 上注意 batch |
|
||
| `torchvision` | `mobilenet_v3_small` / `_large` | 576 / 960 | 面向实时部署 |
|
||
| `torchvision` | `shufflenet_v2_x1_0` | 1024 | 同上 |
|
||
|
||
全是纯卷积(ViT/Swin 不在支持范围,本项目要部署在实时设备上)。改法是**整块替换**
|
||
`simclr.backbone`,因为两种骨架的字段完全不同(`widths/depths/block` vs `arch/pretrained`):
|
||
|
||
```yaml
|
||
# configs/cluster/film.yaml
|
||
simclr:
|
||
backbone:
|
||
name: torchvision
|
||
arch: resnet18
|
||
pretrained: true # 先用 ImageNet 权重起步,再在本批数据上自监督微调
|
||
in_channels: 3
|
||
```
|
||
|
||
如果不想改原文件,用 `_base_` 继承时要加 `_replace_: true`——子字典默认是**合并**而不是
|
||
替换,不加的话 base 里 `tiny_conv` 的 `widths/depths` 会残留下来传给新骨架:
|
||
|
||
```yaml
|
||
_base_: [configs/cluster/film.yaml]
|
||
simclr:
|
||
backbone:
|
||
_replace_: true # 整块换掉,不与 base 合并
|
||
name: torchvision
|
||
arch: resnet18
|
||
pretrained: true
|
||
```
|
||
|
||
> **不要用 `--set simclr.backbone.name=torchvision`**:它只改名字,旧骨架的专有字段仍在,
|
||
> 会报 `构建 'torchvision' 时收到它不接受的字段 ['block', 'depths', 'widths']`。
|
||
> 换组件请改配置文件。
|
||
|
||
注意 `simclr.feature_stages`:`torchvision` 系列的 `features` 是 `nn.Sequential`,可以填
|
||
`[5, 6]` 这类中间层下标;`tiny_conv` 不是,必须保持 `null`(取最终输出)。
|
||
|
||
SimCLR 训练同样接了 SwanLab(配置在 `simclr.swanlab`,字段与 `train.swanlab` 相同,
|
||
默认 `mode: local` 不需要登录)。训练跑起来后另开一个终端:
|
||
|
||
```bash
|
||
uv run swanlab watch runs/simclr/swanlog
|
||
```
|
||
|
||
记录的指标:按 step 的 `simclr/loss`、`simclr/loss_avg`、`simclr/lr`,按 epoch 的
|
||
`simclr/epoch_loss`、`simclr/epoch_seconds`、`simclr/gpu_mem_peak_mb`,超参是完整的 yaml 配置。
|
||
不想要就 `--set simclr.swanlab.enabled=false`;纯看文本日志用
|
||
`grep "完成 loss" runs/simclr/cluster.log`。
|
||
|
||
训练出的骨架和检测模型用的是同一套卷积骨架,可以直接拿去初始化 `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 . # 静态检查
|
||
```
|