Files

869 lines
46 KiB
Markdown
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# ModelTrainEval
薄膜缺陷的**无监督归类 → 分类模型 → 弱监督定位**流水线。模型最终部署在实时设备上,
因此只使用基于 convolution 的结构(不含 attention / transformer)。
完整工作流五步:
1. **自监督**:在无标注数据上用 SimCLR 训练卷积骨架
2. **聚类**:按指定簇数分组,产出可人工翻看的硬链接目录
3. **人工归类**:在 `runs/` 下直接整理簇(可对某一类继续聚类细化),再导出成数据集
4. **分类训练**:在归类好的数据集上训练分类模型
5. **定位**:用 Grad-CAM / SHAP 生成热力图与 YOLO 框标签
## 目录
| 目录 | 作用 |
| --- | --- |
| `configs/` | yaml 配置(`cluster/` 聚类与自监督、`classify/` 分类、`detect/` 检测) |
| `models/` | 模型骨架:`backbones/` + `necks/` + `heads/`,由 `defect_net.py` 组装 |
| `datasets/` | 抛出数据样本的 Dataset 类、数据增强、监督信号编码 |
| `data/` | 数据集本体或软链接(不入库,见 [data/README.md](data/README.md)) |
| `scripts/` | 聚类、框生成、数据集装配与切片、OpenVINO 导出等批处理脚本 |
| `engine/` | 训练 / 评估循环、损失、指标、优化器、检查点、SwanLab 可视化 |
| `utils/` | 配置、注册表、日志、随机种子、设备解析(cuda/mps/cpu) |
| `tests/` | 单元测试 |
| `docs/` | 算法与数学说明文档 |
| `figures/` | 文档用的 mermaid 图(`.mmd` 源 + `.md` 预览 + `.png`) |
入口:[classify.py](classify.py)(分类训练与评估)、[detect.py](detect.py)(YOLO 检测训练与评估),
以及 [scripts/](scripts/) 下的聚类、框生成、导出脚本。
## 环境
```bash
uv sync # 安装依赖(含 swanlab)
uv run python scripts/check_env.py # 自检(Python / PyTorch / 可用后端)
```
### 计算后端
支持 **CUDA / MPS(Apple Silicon)/ CPU**。配置里的 `device: auto` 按 `cuda > mps > cpu`
自动挑最快的可用后端(`classify.py` / `detect.py` / `cluster_unlabeled.py` 都一样),也可以显式写
`cuda` / `mps` / `cpu`,或用命令行 `--device` 覆盖:
```bash
uv run python classify.py --config configs/classify/film.yaml # auto
uv run python classify.py --config configs/classify/film.yaml --device mps # 强制 Mac GPU
uv run python classify.py --config configs/classify/film.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 侧下发耗时。
## 训练可视化(SwanLab)
训练过程通过 [SwanLab](https://swanlab.cn) 记录:按 step 记录各项损失与学习率,按 epoch 记录
验证指标与显存峰值,实验的超参就是完整的 yaml 配置。分类(`configs/classify/`)、检测
(`configs/detect/`)与自监督(`simclr.swanlab`)都在各自配置的 `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 classify.py --config configs/classify/film.yaml --set train.swanlab.mode=local
uv run swanlab watch runs/classify/swanlog # 默认 http://127.0.0.1:5092
# 云端
uv run swanlab login # 首次使用需登录
uv run python classify.py --config configs/classify/film.yaml # mode 默认 online
# 离线跑完后再补传云端
uv run swanlab sync runs/classify/swanlog
```
看板可以在训练**进行中**打开,指标边训边刷。端口冲突用 `-p 5100`,要局域网访问加 `-h 0.0.0.0`。
本地看板依赖 `swanlab[dashboard]`,已包含在 `uv sync` 里。
未安装 swanlab、未登录或网络不通时会打一条警告并继续训练,不会中断。关掉用
`--set train.swanlab.enabled=false`。
## 未分类样本的无监督聚类
### 调用 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 张是离质心最远的离群样本——最能暴露这个簇是不是混了东西。
`clusters/` 里放的是**硬链接**(`link_mode: hard`,默认):不占额外磁盘,但表现得和真实
文件一样——看图软件能开、可以直接在 `runs/` 下把图挪来挪去、建子目录继续分子类,
整个目录拷到别的机器也不会断链。要求与源图同一文件系统,跨盘会自动退回软链并警告。
| `link_mode` | 说明 |
| --- | --- |
| `hard`(默认) | 硬链接。零额外磁盘,行为同真实文件,可跨机器拷贝 |
| `sym` | 软链接。跨盘可用,但记绝对路径,换机器就断(`export_clusters.py` 能重定位) |
| `copy` | 直接复制。跨盘可用且绝不断链,代价是占一份磁盘 |
> `hard` 与原图是**同一份数据**:在 `clusters/` 里改图会改到原图;只挪动、删除条目
> 不影响原图。
**尺寸不统一的处理**:这批图是已裁好的缺陷小图,中位 152×113,宽高比从 0.24 到 39.9。
尺寸和宽高比是上游 ROI 裁剪的产物、不是缺陷属性,所以预处理的目标是把它抹掉:
直接拉伸到 128×128(不用 letterbox——填充比例等价于把宽高比明着喂给网络),
再做逐图 z-score 抹掉整体亮度/对比度。
无标签时用两个指标判断质量:**轮廓系数**(簇是否紧致)和**换随机种子的稳定性**
(同一份数据换初始化后划分是否还一样,低说明是在硬切一团连续数据)。
> 完整的数学推导、算法流程图与全部消融数据见
> **[docs/unsupervised-clustering.md](docs/unsupervised-clustering.md)**。
### 按聚类结果导出数据集
聚类产出的 `clusters/` 默认是**硬链接**(见上文 `link_mode`),方便人工翻看和调整——
可以把图在子目录间挪动、删掉分错的、把 `cluster_03` 改名成 `grease`。整理完之后,
以这个目录为准把图**复制**出来,得到一个与 `runs/` 彻底解耦的数据集:
```bash
uv run python scripts/export_clusters.py --work-dir runs/simclr_r101 --prefix-parent
# 先看计划加 --dry-run
```
```
data/simclr_r101/ # 默认取 work_dir 的名字,--out 可完全指定
├── grease/ # 子目录名沿用 clusters/ 里的(人工整理过的类别名)
├── mosquito/
├── normal/
└── manifest.json # 每张图的原始路径与所属类别,便于追溯
```
导出的是**独立的副本**:硬链接条目虽然本身就是原图数据,但路径在 `runs/` 下、
文件名还带排序前缀,而且删改会牵连原图。文件名用原图的名字、去掉 `clusters/` 里的
`0125_` 排序前缀。加 `--prefix-parent` 会再拼上原父目录名(`反面纵黑点/abc.png` ->
`反面纵黑点_abc.png`),既保住来源信息,也化解不同目录同名文件撞车。
> **三种条目都能解析。** 硬链接(默认)条目本身就是原图,脚本按 inode 反查
> `assignments.csv` 认出它对应哪张原图,好让文件名与 `--prefix-parent` 拿到真正的来源;
> 有效软链直接 resolve;**从服务器拷回来的软链**(`link_mode: sym` 时,指向
> `/home/deep/...`)在本机全部失效,脚本读出链接目标、按其中的 `data/` 段重新落到本机
> `--source-root`(默认当前目录)上,万一还找不到再用文件名在源目录里兜底索引一次。
> 解析不到的条目会列出来并中止,确认要忽略就加 `--skip-missing`。
几个保护:同类别内重名(保持原名会互相覆盖)直接报错并提示可用 `--prefix-parent`;
目标目录非空时必须显式 `--overwrite`。
### 自监督预训练(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。
**续训到调度末段时想调大 lr**:cosine 是 epoch 的函数,在 2229/2500 处只剩满值的 3%,
所以新设的 `lr=1e-3` 实际只有 `2.9e-5`。日志会警告这一点。要让新 lr 真正生效:
```bash
# 按剩余轮数重开一条 cosine 曲线,新 lr 从满值开始
--set simclr.restart_schedule=true
# 或者拉长整条曲线(总轮数变大,当前位置的系数就回升)
--set simclr.epochs=5000
```
续训时各项配置的行为:
| 配置项 | 续训时 |
| --- | --- |
| `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` 会报错,因为它不含投影头和优化器状态。
#### 训练与聚类分开跑
自监督训练和聚类在代码上本来就是解耦的(`engine/simclr.py` 与 `engine/cluster.py`
互不引用,唯一接口是 `simclr_backbone.pth`),只是默认一条命令串起来跑。
需要分两步时——比如在 GPU 机器上专心训练几十小时,之后在别处反复调 k——用
`--train-only` 断开:
```bash
# 第一步:只训练,存完权重就退出,不抽特征也不聚类
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film_simclr_resnet18.yaml \
--work-dir runs/simclr_r18 --train-only --set simclr.epochs=40
# 第二步:只聚类,想换几次 k 就跑几次,每次约 1 分钟
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film_simclr_resnet18.yaml \
--work-dir runs/simclr_r18_k4 \
--set simclr.checkpoint=runs/simclr_r18/simclr_backbone.pth --num-clusters 4
```
第一步只产出 `simclr_backbone.pth` 与 `simclr_last.pth`,没有 `assignments.csv` 之类。
`--train-only` 遇到 `method=pretrained`(没东西可训)、`simclr.checkpoint`(会跳过训练)
或 `--sweep-k`(要靠聚类才能扫)都会直接报错,不会跑完才发现什么都没产出。
#### 第二步:复用权重重新聚类
想换簇数或改聚类参数时**不必重训**——把权重路径填进 `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 / resnet101 的现成配置**,直接用:
```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
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film_simclr_resnet101.yaml \
--work-dir runs/simclr_r101 --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 |
| resnet101 | 42.5M | 1536 | 未实测 | 配置里用 128(按规模比例保守取值) |
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` | `resnet101` | 2048 | 最深,首次使用会下载约 171 MB 权重 |
| `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 权重更贴合薄膜的成像特性。
## 缺陷分类模型
在聚类 + 人工归类得到的数据集上训练分类模型(`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` 可以单独衡量
自监督特征本身有多好。
## 从分类到定位:生成框标签
分类模型只学过"这是什么缺陷",没学过"缺陷在哪"。归因方法能把分类分数反推回空间位置,
阈值化 + 取连通域外接矩形,就得到弱监督的框标签。另外还有一条不依赖模型的传统 CV 路线:
```bash
# Grad-CAM(默认,全量 11k 张约 2 分钟)
uv run python scripts/generate_boxes.py --out runs/boxes_gradcam --method gradcam
# SHAP Partition(更精细,全量约 1.7 小时)
uv run python scripts/generate_boxes.py --out runs/boxes_shap --method shap
# 传统 CV 背景差分(最快,且不需要分类模型)
uv run python scripts/generate_boxes.py --out runs/boxes_cv --method cv
```
**默认对全量数据生成**;`--limit` / `--per-class` 只用于快速试参数。
### 产出是自成一体的 YOLO 数据集
```
runs/<name>/
├── images/<类别>/<图名>.png # 硬链接到原图,不占额外磁盘
├── labels/<类别>/<图名>.txt # YOLO:class_id cx cy w h(归一化)
├── heatmaps/<类别>/<图名>.png # 热力图叠加原图 + 框,供人工核对
├── data.yaml
└── summary.json
```
`images/` 与 `labels/` 的对应关系符合本仓库的 YOLO 约定,所以 `FilmDefectDataset`
能直接读——这一步把分类的产出接回了检测训练链路。
热力图与框的类别默认取**图片所在目录的人工标注**(`--target label`)。类别的真值是人工归类,
分类模型只用来定位;改用 `--target predicted`(模型 top1 预测)时,模型分错的图框的类别也会跟着错。
两者只在预测与目录不一致的图上有区别——`boxes_shap` 那次是 0.85%,约 95 张。
`summary.json` 里的 `prediction_matches_folder_label` 记录两者的一致率(实测 0.987)。
### 三种方法的取舍
| | `gradcam` | `shap` | `cv` |
| --- | --- | --- | --- |
| 依赖分类模型 | 是 | 是 | **否** |
| 分辨率 | 受限于特征图(layer3 @128² 只有 8×8) | 与输入同分辨率(128×128) | 原图分辨率 |
| 单张耗时 | ~0.01 秒 | ~0.55 秒 | ~0.011 秒 |
| 小目标 | 几像素的黑点在 8×8 上不可分辨 | 能贴合到几十像素的目标 | 贴合好 |
| 找的是什么 | 模型据以判断类别的区域 | 同左,更精细 | 偏离局部背景的区域 |
实测框面积占比(越小越贴合缺陷本体):`blackdot` 0.019(shap) vs 0.269(gradcam)、
`mosquito` 0.030 vs 0.173、`whitedot` 0.076 vs 0.253。小目标类别上 SHAP 明显更准,
代价是慢 50 倍。SHAP 需要 `uv sync --group explain`。
### 传统 CV:背景差分(`--method cv`)
实现在 [engine/classical.py](engine/classical.py):大半径方框均值估背景 → 残差(可选只看
比背景暗/亮的一侧)→ 小半径平滑 → 判定有无缺陷 → 按峰值归一化,最后交给和另外两种方法
**同一套** `cam_to_boxes` 出框。只依赖 numpy。
它与归因方法有三点本质差别:
- **不需要分类模型**,`--checkpoint` 都不用给,分类器还没训练出来就能先跑一版框;
也因此不受分类精度影响。
- **类别只能来自目录**(`--target label`),它自己不判类。
- **找的是"看起来和周围不一样"的区域**,不是"模型据以判断类别的区域"。对
`blackdot` / `whitedot` / `mosquito` 这类局部目标效果好;`ripple` / `wrinkle` 这种铺满整图的
纹理类本来就没有局部目标,它给不出有意义的框。
**有无缺陷是显式判定的**,这点和归因方法不同。残差峰值要同时高于绝对对比度下限
(`--cv-min-contrast`,默认 0.05≈13 个灰阶)和噪声水平的 15 倍,否则直接返回全零图、出 0 个框。
噪声水平用残差的 MAD 估计。这条判定是必需的:实测纯噪声图的"残差峰值 / MAD"稳定在 9 附近
(最大 10.9,几乎与噪声大小无关),而有一个对比度 0.25 的黑点时是 25~330,两者分得很开;
若只按峰值归一化,一张正常图的噪声会被拉到 1,满图假框。所以正常图天然得到空标签。
调参:`--cv-radius`(背景半径,0=按图取 `min(高,宽)//8`;太小会把缺陷自身算进背景)、
`--cv-polarity`(`dark` 黑点蚊虫 / `bright` 白点反光 / `both`)、`--cv-min-contrast`、
`--smooth-radius`(与 shap 共用)、`--cv-max-side`(分析前缩边长,默认 256——框是归一化的,
缩放不影响标签,但能挡住纯 Python 的连通域标记在大图上变慢)。
> **已知局限。** 背景估计在图像边界处会对强烈的明暗阶跃产生响应,所以贴着边框的
> 明暗交界可能被框出来。对 `edge_clean` / `edge_dirty` 这种本来就在拍边缘的类别,这是真实信号;
> 对其他类别则是假框,复核时留意。
### 调参
- `--max-boxes`(默认 1):一张裁片通常只对应一处缺陷,多出来的框多半是噪声
- `--threshold`(默认 0.4):热力图相对最大值的阈值
- `--min-area-ratio`:**缺省随方法而变**(gradcam 0.02、shap 0.002)。两者归因图分辨率
差 16 倍,同一个小黑点在 128×128 上只占百来像素,沿用 0.02(=328 像素)会被整个滤掉
- `--peak-ratio`(默认 0.75):次要连通域的峰值需达全局最大值的这个比例,滤掉弱响应假框
- `--padding`(默认 0.1):框四周外扩。归因倾向只覆盖最具判别性的部位,往往偏小
- `--stage`(gradcam,默认 6=layer3):layer4 只有 4×4 太粗,layer2 响应太散
- `--no-heatmaps` / `--no-images`:只要标签时可跳过
> **不是所有类别都可定位。** `normal`/`black`/`edge_clean` 描述的是整图属性而非局部缺陷,
> 它们的框没有物理意义。真正有用的是 `blackdot`/`whitedot`/`mosquito`/`grease` 这些
> 有局部目标的类别。
## 检测模型训练(YOLO)
`generate_boxes.py` 产出的是弱监督框,要先在 Label Studio 里人工修正,再划分 train/val 训练:
```bash
# 1) 生成 Label Studio 复核材料到 runs/boxes_shap_label_studio/
uv run python scripts/prepare_label_studio.py --source runs/boxes_shap
# 先看计划加 --dry-run
# 2) 在 Label Studio 里修正标签(导入步骤见下文,脚本结束时也会打印)
# 3) 导出修好的标签,划分 train/val 到 data/boxes_shap/(导出脚本待写)
# 4) 训练
uv run python detect.py --config configs/detect/film.yaml
# 只评估 / 续训
uv run python detect.py --config configs/detect/film.yaml --eval-only \
--checkpoint runs/detect/film/weights/best.pt
uv run python detect.py --config configs/detect/film.yaml --resume
```
### 准备 Label Studio 复核材料
```
runs/boxes_shap_label_studio/ # 同时也是 Label Studio 的 DOCUMENT_ROOT
├── images/<类别>_<序号>.png # 一律 PNG:PNG 原图硬链接,BMP 等无损转换
├── labels/<类别>_<序号>.txt
├── classes.txt # 每行一个类别名,行号即 class_id
├── name_map.csv # 新文件名 ↔ 源文件路径
├── summary.json # 逐类数量、空标签处置、落盘方式
├── tasks.json # 每张图一个任务,框作为 predictions
└── tasks.label_config.xml # 各类别的 RectangleLabels
```
**不写 `data.yaml`、不划分 train/val。** 这是给人复核的材料,不是训练集,没有 `data.yaml`
也就不会被误拿去训练。复核期间 Label Studio 直接从这个目录读图,**不要删除或 `--overwrite` 它**。
`detect.py` 若检测到 `data.yaml` 的 train 与 val 相同(`generate_boxes.py` 的产出就是这样),
也会告警:验证指标算的是训练集本身,mAP 虚高,早停与 `best.pt` 的选择失真。
**图片统一成 PNG。** pearlescent / blue / uv 的原始数据是 BMP(共 10,814 张),
`generate_boxes.py` 按原后缀链接,于是产出里混着两种格式。这一步按**文件头**
(不是后缀)判断:真 PNG 照旧硬链接不占磁盘,其余无损转 PNG——不缩放像素,归一化的框坐标
原样有效;CMYK 这类 PNG 存不下的模式转 RGB。转换无法硬链接,会占额外磁盘,`--dry-run`
会先报出要转多少张。统一后缀后 `a.png` 与另一目录的 `a.bmp` 会撞名,脚本直接报错。
**按类别重新编号。** 源文件名五花八门(`反面纵黑点_1_550mmX1561.71m_20260416160441_0.png`、
`pearlescent_404.bmp`……),输出统一改成 `<类别>_<序号>.png`,如 `blackdot_0042.png`。
类别以标签第一列为准,文件名永远不会与标签矛盾;序号在类内按源路径排序,同一份来源重跑,
同一张图拿到的名字不变。原文件名里的正反面、位置、时间戳等信息都在 `name_map.csv` 里
(带 BOM,Excel 直接打开中文不乱码)。想保留原名加 `--keep-names`。
**classes.txt** 是 label-studio-converter 读类别用的(也兼容 LabelImg / X-AnyLabeling),
行号就是 class_id,内容与来源 `data.yaml` 严格一致。
### 导入 Label Studio
**图片 URL 必须按 Label Studio 实际的 `DOCUMENT_ROOT` 生成。** 任务里写的是
`/data/local-files/?d=<相对路径>/images/<图名>.png`,Label Studio 把 `d` 拼到
`LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT` 后面找文件,相对路径错了就是整个项目 404。
脚本缺省读这个环境变量,也可以用 `--document-root` 显式指定:
```bash
# DOCUMENT_ROOT 统一放标注项目,产出仍留在 runs/
uv run python scripts/prepare_label_studio.py --source runs/boxes_shap \
--document-root /home/deep/labeling
```
| 情形 | 做法 | 图片 URL |
|---|---|---|
| 没设 DOCUMENT_ROOT | 输出目录自己当 DOCUMENT_ROOT | `?d=images/…` |
| 输出目录在 DOCUMENT_ROOT 之下 | 直接用 | `?d=<相对路径>/images/…` |
| 输出目录在 DOCUMENT_ROOT 之外 | 在 DOCUMENT_ROOT 里建同名软链接指向它 | `?d=boxes_shap_label_studio/images/…` |
DOCUMENT_ROOT 里已有同名的**真目录**(比如手动复制过去的一份)时脚本会报错,不会覆盖它。
Local files 存储的绝对路径填 `<DOCUMENT_ROOT>/<访问路径>/images`,脚本打印的步骤里已经算好。
脚本结束时会打印带实际路径的导入步骤:设置 `LABEL_STUDIO_LOCAL_FILES_*` 环境变量启动、
粘贴 labeling config、添加 Local files 存储(**只保存不要 Sync**,否则每张图多出一个空任务)、
导入 `tasks.json`。
几个取舍:
- **为什么是扁平结构**:`label-studio-converter import yolo` 只认 `<input>/{images,labels,classes.txt}`,
且 `os.listdir` 不递归,看不见子目录。
- **软链接可以用**:Label Studio(1.23.0 核对过源码)拼文件路径用 `safe_join`(`abspath`),
存储路径做 `normpath` 后按字符串前缀比较,都**不解析软链接**,所以 DOCUMENT_ROOT 里放一个
指向 `runs/` 的软链接就能访问,不必复制。存储路径也必须是 DOCUMENT_ROOT 的**子目录**,
与 DOCUMENT_ROOT 相同会被拒绝,这正是要填到 `…/images` 的原因。
- **框作为 predictions 而不是 annotations**:它们是模型生成的。审核员修改提交后才成为 annotation,
哪些审过一目了然,导出时也只导出 annotation,没审过的框不会混进去。
- **不传 `--image-dims`**:图片尺寸有 3,223 种,传了会把同一个宽高写进每条标注。
- converter 自带的导入提示漏了环境变量(它在循环里改写了 `image_root_url`,与默认值比较失败),
所以它的输出被静音,改由脚本打印。
> **复核重点是 `normal`。** SHAP 在每张图上都框了东西,`normal` 有 4,295 个框、占 38%。
> 若 `normal` 表示无缺陷,这些框要删掉,否则检测器会学着"检测正常"。
### 没有框的图怎么处理
Grad-CAM 没框出东西的有 339 张(占 3%)。它们**默认保留为负样本**,对压低误检有用。
但要当心:**空标签在 YOLO 里的语义是"这张图什么都没有"**——一个明确的负样本断言,
而不是"有缺陷但没能定位出来"。而这 339 张里只有 69 张真的来自 `normal`,其余 270 张是人工
标成缺陷、Grad-CAM 却失手的图:
| 来源类别 | 空标签 | 占该类 |
|---|---|---|
| grease | 95 | 7.8% |
| blackdot | 82 | 11.3% |
| normal | 69 | 2.6% |
| strip_h | 28 | 2.9% |
| zigzag | 23 | 2.4% |
| black | 17 | 24.3% |
| 其余 9 类 | 25 | <3% |
照搬的话,`blackdot` 有 11% 的正样本、`black` 有 24% 会以"此图无缺陷"的身份进入训练,
等于在教模型抑制真实缺陷。用 `--empty-policy` 按来源类别分别处置:
```bash
uv run python scripts/prepare_label_studio.py --source runs/boxes_gradcam \
--empty-policy drop normal=keep black=full edge_clean=full
```
- `keep`(默认)——当负样本保留。只有真正无缺陷的类别该用
- `drop`——丢弃。缺陷类没定位出来的属于**信息缺失**,不该冒充负样本
- `full`——退化成整图框。适合 `black`/`edge_clean` 这种描述整图属性、缺陷本就铺满全图的类别
第一个裸词是全局默认,后面 `类别=策略` 逐类覆盖。类别名或策略写错会直接报错而不是静默忽略。
没有显式指定策略、又落到 `keep` 的类别会打一条告警,提醒你确认它是否真的无缺陷。
拍平之后不同类别目录下的同名文件会互相覆盖,脚本会直接报错而不是静默丢数据;
真遇上就用 `scripts/export_clusters.py --prefix-parent` 重新导出。
### 模型与增强的取舍
- **默认 YOLO26 的 `n` 档**:CLAUDE.md 的"只用 convolution 模型"只约束聚类与数据集生成阶段,
检测这一步不受限,所以用最新的 YOLO26——它是端到端的(NMS-free),推理侧少一道后处理。
选最小的 `n`:数据只有 1.1 万张且标签是弱监督的,模型再大容易把标注噪声一起学进去。
- **关掉 mosaic / mixup**:它们把四张图拼一起,而我们每张图只有一个弱监督框,拼接后小目标被
缩得更小,对 `blackdot`/`whitedot` 这类本就只占千分之几面积的类别是净伤害。
- **关掉 hsv_h**:薄膜图近乎单色,色相扰动没有意义。翻转照旧开着(缺陷没有固定朝向)。
- 冒烟测试用 `--set train.epochs=1 train.fraction=0.02 train.imgsz=320` 先把链路跑通。
> **指标怎么读。** 框来自 Grad-CAM/SHAP 的弱监督标注,不是人工真值。mAP 衡量的是"检测器
> 能不能复现归因热力区域",不等于"缺陷定位有多准"。真值上的精度仍需人工抽检。
### 可视化
`train.swanlab.enabled: true` 时接 SwanLab 回调,训练曲线实时可看(`swanlab watch` 开本地看板)。
ultralytics 自己还会在 `runs/detect/film/` 下写:
- `results.png` / `results.csv`——各项 loss 与 mAP 曲线
- `confusion_matrix.png`、`confusion_matrix_normalized.png`
- `BoxPR_curve.png` / `BoxP_curve.png` / `BoxR_curve.png` / `BoxF1_curve.png`
- `train_batch*.jpg`、`val_batch*_labels.jpg` 与 `val_batch*_pred.jpg`——标注与预测的直观对照
- `metrics.json`——mAP50-95 / mAP50 / P / R 与逐类 mAP50
> **许可提醒。** ultralytics 是 AGPL-3.0。模型若要随产品分发,需要开源整套代码或另购商业授权;
> 仅内部使用不受影响。
## 方法
- **骨架**:`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` 引用。
## 分类模型导出 OpenVINO
分类器直接输出按训练类别顺序排列的 Softmax 概率,不适配 YOLO 检测输出:
```bash
uv pip install openvino
uv run python scripts/export_classifier_openvino.py \
--checkpoint runs/classify/best.pth \
--fp32-weights --out runs/classify/models/resnet18_128_probabilities
```
脚本读取分类检查点的配置和同目录 `eval_report.json` 中的类别顺序;也可用
`--classes classes.txt` 或 `--classes eval_report.json` 明确指定分类器的类别顺序。
指定顺序不能与检查点记录矛盾。输入尺寸从检查点读取,不依赖 YOLO 模型。
产出的 XML/BIN 按骨架命名,如 `resnet18.xml/.bin`、`mobilenet_v3_small.xml/.bin`。
输入为 RGB float32 BCHW、值域 `[0,1]`,输入名为 `images`。
逐图标准化仅在训练配置使用时嵌入模型。脚本不进行 resize,图片必须已切为目标尺寸。
输出名为 `probabilities`,形状 `[1,C]`:九类模型是 `[1,9]`,索引直接对应 `classes.txt`。
每类概率在 `[0,1]` 内,概率和为 1;normal 也返回真实概率,不清零、不做阈值过滤。
部署端直接对 C 个概率取 argmax;normal 是否表示正常、是否应用阈值由业务层决定。
现有 `deploy/cpp` 是 YOLO 框解码器,不能用于此分类输出;需要按 `[1,C]` 读取,
不要跳过前四个通道或遍历候选框。模型附带 `metadata.yaml`、`classes.txt` 和 `export.json`。
输出目录必须是新目录,原有模型不会被覆盖。
默认压缩权重为 FP16;需要优先核对转换前后判别一致性时,加 `--fp32-weights`
保留 FP32 权重,减少分类边界处的舍入差异。输入输出接口仍是 float32。
### 在固定 128×128 YOLO 裁片上监督训练分类器
```bash
uv run python classify.py --config configs/classify/film_tiles128_resnet18.yaml \
--work-dir runs/classify_tiles128_resnet18
uv run python classify.py --config configs/classify/film_tiles128_mobilenet_v3_small.yaml \
--work-dir runs/classify_tiles128_mobilenet_v3_small
```
数据从 `data/film_yolo_128/data.yaml` 读取,保留现有 train/val 划分和类别编号。
监督标签来自 YOLO 标注文件,空标签表示 normal;多类别裁片取面积最大的框对应类别。
缺失标签文件会报错,不会自动当背景。图片必须已经是 128×128 PNG,禁止 resize。
这两份配置使用 RGB / 255,与检测部署输入一致;不进行逐图标准化。
最佳权重按验证 macro-F1 保存,可用上述导出脚本生成 OpenVINO 分类概率模型。
分类准确率、转换前后一致率和与 YOLO 的一致率是不同指标,应分别报告。
导出后,在同一固定裁片验证集上检查这三个指标(示例):
```bash
uv run python scripts/compare_openvino_classes.py \
--teacher runs/detect/models/film_128_20261013_0902_openvino_model/film_128_20261013_0902.xml \
--student runs/classify/models/resnet18_tiles128/resnet18.xml \
--checkpoint runs/classify_tiles128_resnet18/best.pth \
--conf 0.25 --report runs/classify_tiles128_resnet18/comparison.json
```
`--student` 指向实际导出的 XML。这里 `--teacher` 只用于比较,不参与训练。
YOLO 有有效框时取最高分候选的类别,无有效框记为 normal。比较工具在业务层对分类概率
使用同一置信度阈值,低于阈值时记为 normal;模型本身仍返回全部真实概率。
报告同时记录各类支持数、总体/前景一致率和逐图结果。
## 开发
```bash
uv run pytest # 单元测试
uv run ruff check . # 静态检查
```