Files
ModelTrainEval/README.md
T
Shaun ZhouandClaude Opus 5 0f5bb4a5c0 feat(cluster): 加 --train-only,自监督训练与聚类可分两步跑
engine/simclr.py 与 engine/cluster.py 本来就互不引用,唯一接口是
simclr_backbone.pth;耦合只在 CLI——method=simclr 会训练完直接聚类。
此前"只聚类"可以用 simclr.checkpoint 做到,"只训练"没有办法。

- --train-only:训练存完权重即退出,不抽特征也不聚类,并打印下一步的聚类命令
- 与 method=pretrained / simclr.checkpoint / --sweep-k 组合时立刻报错说明原因,
  而不是跑完才发现没有产出
- 典型用法:GPU 机器上 --train-only 跑几十小时,之后拿权重反复调 k,每次约 1 分钟

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-07 11:38:59 +08:00

556 lines
27 KiB
Markdown
Raw 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),并支持导出 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` 会报错,因为它不含投影头和优化器状态。
#### 训练与聚类分开跑
自监督训练和聚类在代码上本来就是解耦的(`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 的现成配置**,直接用:
```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 . # 静态检查
```