diff --git a/README.md b/README.md index 5e5398c..344b18d 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,7 @@ convolution 的结构(不含 attention / transformer),并支持导出 ONNX | `data/` | 数据集本体或软链接(不入库,见 [data/README.md](data/README.md)) | | `scripts/` | 格式转换、数据集切分、ONNX 导出等批处理脚本 | | `engine/` | 训练 / 评估循环、损失、指标、优化器、检查点、SwanLab 可视化 | -| `utils/` | 配置、注册表、日志、随机种子、GPU 守卫 | +| `utils/` | 配置、注册表、日志、随机种子、设备解析(cuda/mps/cpu) | | `tests/` | 单元测试 | | `docs/` | 算法与数学说明文档 | | `figures/` | 文档用的 mermaid 图(`.mmd` 源 + `.md` 预览 + `.png`) | @@ -148,10 +148,12 @@ 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 内存)。 +输出包含:每个 batch 的单帧延迟 / FPS、显存峰值、权重占用,以及按显存和耗时各排一次序的 +算子级明细。`--amp` 用 fp16 推理对比,`--postprocess` 把 heatmap 解码也计入, +`--device` 指定后端(默认 auto)。 + +后端差异见[计算后端](#计算后端):CUDA 给精确峰值;MPS 上 `torch.profiler` 抓不到 Metal kernel, +算子明细只反映 CPU 侧下发耗时、显存靠采样是下界;CPU 只有耗时与 CPU 内存。 ## 未分类样本的无监督聚类 @@ -195,27 +197,71 @@ uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml -- 视图互为正样本,batch 内其余 2B−2 个视图为负样本,用 NT-Xent 损失训练。 训练走 `device: auto` 选中的后端——**Mac 上直接用 MPS 即可**,实测 14,579 张、 -batch 256 时约 65 秒/epoch: +batch 256 时约 65 秒/epoch。 + +#### 第一步:训练 + 聚类(一条命令) + +`method=simclr` 会先训练、存权重,然后直接用这份权重抽特征做聚类,不需要分两次调用: ```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 +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 # 每 save_interval 个 epoch 覆盖一次;Ctrl-C 也会先落盘 +├── 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 +``` + +#### 第二步:复用权重重新聚类 + +想换簇数或改聚类参数时**不必重训**——把权重路径填进 `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-预训练)。 -训练每 `save_interval` 个 epoch 落一次盘,中途 Ctrl-C 也能拿到当前权重继续用。 权重存到 `/simclr_backbone.pth`(**只含骨架,不含投影头**——投影头被训练成 -对增强不变,会丢掉下游有用的信息,所以聚类和保存都只用骨架特征)。相关配置: +对增强不变,会丢掉下游有用的信息,所以聚类和保存都只用骨架特征)。 + +#### 固化到配置 + +每次敲 `--set` 麻烦的话,直接改 [configs/cluster/film.yaml](configs/cluster/film.yaml): ```yaml method: simclr @@ -224,12 +270,12 @@ simclr: # 自监督用项目自己的纯卷积骨架(与 method=pretrained 的顶层 `backbone` 各自独立, # 两者字段完全不同,分开配置以免 _base_ 合并时串味) backbone: {name: tiny_conv, in_channels: 3, widths: [32, 64, 128, 256], depths: [1, 2, 2, 2]} - epochs: 100 + epochs: 40 batch_size: 256 temperature: 0.2 # 越小越关注难负样本 amp: auto # CUDA 上开、MPS 上关 save_interval: 5 # 每 5 个 epoch 落一次盘 - checkpoint: null # 填已有权重可跳过重复训练 + checkpoint: null # 填 runs/simclr/simclr_backbone.pth 即可跳过重复训练 projection: {hidden_channels: 256, out_channels: 128} optimizer: {name: adamw, lr: 0.001, weight_decay: 0.0001} scheduler: {name: cosine} @@ -244,6 +290,9 @@ view_transforms: # 生成两个视图的随机增强 > `RandomResizedCrop` 的 `scale` 下限别低于 0.4:裁太狠会让两个视图都只剩背景, > "背景 vs 背景"被当成正样本对,学出来的特征会退化。 +> SimCLR 训练目前**没有接 SwanLab**,只往 `/cluster.log` 打文本日志, +> 盯 loss 要看日志(`grep "完成 loss" runs/simclr/cluster.log`)。 + 训练出的骨架和检测模型用的是同一套卷积骨架,可以直接拿去初始化 `DefectNet` 的 backbone—— 标注量有限时,通常比 ImageNet 权重更贴合薄膜的成像特性。