ModelTrainEval

薄膜缺陷分类 + 定位模型的训练与评估。模型最终部署在实时设备上，因此只使用基于 convolution 的结构（不含 attention / transformer），并支持导出 ONNX。

目录

目录	作用
configs/	yaml 配置（model/ 模型、dataset/ 数据、train/ 训练、eval/ 评估）
models/	模型骨架：backbones/ + necks/ + heads/，由 defect_net.py 组装
datasets/	抛出数据样本的 Dataset 类、数据增强、监督信号编码
data/	数据集本体或软链接（不入库，见 data/README.md）
scripts/	格式转换、数据集切分、ONNX 导出等批处理脚本
engine/	训练 / 评估循环、损失、指标、优化器、检查点、SwanLab 可视化
utils/	配置、注册表、日志、随机种子、GPU 守卫
tests/	单元测试
docs/	算法与数学说明文档
figures/	文档用的 mermaid 图（.mmd 源 + .md 预览 + .png）

入口在仓库根目录：train.py、eval.py。

环境

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

# 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 与 datasets/yolo.py。

训练与评估

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

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 记录：按 step 记录各项损失与学习率，按 epoch 记录 验证指标与显存峰值，实验的超参就是完整的 yaml 配置。在 train.swanlab 里配置：

train:
  swanlab:
    enabled: true
    project: film-defect
    name: null        # 缺省用 work_dir 目录名
    mode: online      # online(云端) | local(本地看板) | offline(仅存盘) | disabled
    tags: [tiny_conv, center-head]

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 看单帧耗时与显存占用：

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 里大量样本还没人工分完。这条流水线把目录下所有图片拉平（不看任何 目录名），做无监督聚类，产出可人工核对的分组：

# 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。

自监督预训练（SimCLR）

method: pretrained（默认）用 ImageNet 权重直接抽特征，不训练；method: simclr 则在 本批薄膜图上自监督训练一个项目自己的纯卷积骨架，域适配更好。同一张图的两个随机增强 视图互为正样本，batch 内其余 2B−2 个视图为负样本，用 NT-Xent 损失训练。

这属于训练，因此必须有 Nvidia GPU（走 utils.device.require_cuda() 守卫， 在无卡机器上会直接报错退出）：

# 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（只含骨架，不含投影头——投影头被训练成 对增强不变，会丢掉下游有用的信息，所以聚类和保存都只用骨架特征）。相关配置：

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。 新增骨架 / 数据集只需用注册表装饰器注册，再在 yaml 里按 name 引用。

开发

uv run pytest         # 单元测试
uv run ruff check .   # 静态检查