diff --git a/README.md b/README.md index 24d79e5..1fde3ad 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,15 @@ # ModelTrainEval -薄膜缺陷**分类 + 定位**模型的训练与评估。模型最终部署在实时设备上，因此只使用基于 -convolution 的结构（不含 attention / transformer），并支持导出 ONNX。 +薄膜缺陷的**无监督归类 → 分类模型 → 弱监督定位**流水线。模型最终部署在实时设备上， +因此只使用基于 convolution 的结构（不含 attention / transformer）。 + +完整工作流五步： + +1. **自监督**：在无标注数据上用 SimCLR 训练卷积骨架 +2. **聚类**：按指定簇数分组，产出可人工翻看的硬链接目录 +3. **人工归类**：在 `runs/` 下直接整理簇（可对某一类继续聚类细化），再导出成数据集 +4. **分类训练**：在归类好的数据集上训练分类模型 +5. **定位**：用 Grad-CAM / SHAP 生成热力图与 YOLO 框标签 ## 目录 @@ -18,7 +26,7 @@ convolution 的结构（不含 attention / transformer），并支持导出 ONNX | `docs/` | 算法与数学说明文档 | | `figures/` | 文档用的 mermaid 图（`.mmd` 源 + `.md` 预览 + `.png`） | -入口在仓库根目录：[train.py](train.py)（检测+定位训练）、[eval.py](eval.py)、[classify.py](classify.py)（分类训练与评估）。 +入口：[classify.py](classify.py)（分类训练与评估）、[scripts/](scripts/) 下的聚类、导出与框生成脚本。 ## 环境 @@ -50,98 +58,6 @@ uv run python train.py --config configs/train/baseline.yaml --device cpu `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//`）：`train.log`、`config.yaml`、`last.pth`、`best.pth`、`swanlog/`。 - -**续训以配置为准。** `--resume` 只从检查点取"进度"——模型权重、优化器动量缓冲、 -已训到第几轮、`best_metric`、`global_step`；`epochs`、`lr`、`weight_decay`、调度器 -一律以当前 yaml 为准，所以可以顺手调超参： - -```bash -uv run python train.py --config configs/train/baseline.yaml --resume runs/baseline/last.pth \ - --set train.epochs=200 train.optimizer.lr=5e-4 -``` - -与检查点不一致时日志会逐项列出（`配置与检查点不一致，按当前配置执行：总轮数 100 -> 200； -optimizer.lr 0.001 -> 0.0005`）。学习率调度按新的总轮数重建并快进到当前 epoch。 -`epochs` 是**总轮数**不是剩余轮数。SimCLR 自监督那边是同样的语义。 - -## 多类 / 单类分类 - -图像级分类的形式由 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 记录 @@ -176,26 +92,6 @@ uv run swanlab sync runs/baseline/swanlog 未安装 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 预训练模型 diff --git a/configs/dataset/film_defect.yaml b/configs/dataset/film_defect.yaml deleted file mode 100644 index c891d5a..0000000 --- a/configs/dataset/film_defect.yaml +++ /dev/null @@ -1,51 +0,0 @@ -# 薄膜缺陷数据集（YOLO 风格标注：labels/*.txt，每行 `class_id cx cy w h`，归一化坐标）。 -# 类别顺序即标注里的 class_id，改动顺序等于改变标签含义，必须重新训练。 -# 也可以写成 data.yaml 的路径，从其中的 names 读类别：classes: data/raw/film/data.yaml -classes: [scratch, bubble, particle, crease, stain] - -# 图像级标签形式；模型的分类头、分类损失、分类指标都跟随它。 -# multilabel（多类）：一张图可同时含多种缺陷，多热标签 + BCE，全 0 表示 OK 片； -# multiclass（单类）：一张图只归一个类别，互斥 softmax + CE， -# 分类头多一维表示"无缺陷"（下标 = len(classes)）， -# 一图出现多种缺陷时取框数最多的类别。 -# 无论取哪种，定位分支始终按 classes 逐类输出 heatmap。 -label_mode: multilabel - -image_size: [512, 512] -gray: false - -data: - train: - dataset: - name: film_defect - img_dir: data/raw/film/images/train - label_dir: null # 缺省按 YOLO 约定取 data/raw/film/labels/train - gray: false - max_objects: 64 - transforms: - - {name: Resize, size: [512, 512]} - - {name: RandomHorizontalFlip, p: 0.5} - - {name: RandomVerticalFlip, p: 0.5} - - {name: ColorJitter, brightness: 0.2, contrast: 0.2} - - {name: ToTensor} - - {name: Normalize, mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225]} - loader: - batch_size: 16 - shuffle: true - num_workers: 8 - drop_last: true - - val: - dataset: - name: film_defect - img_dir: data/raw/film/images/val - label_dir: null - gray: false - transforms: - - {name: Resize, size: [512, 512]} - - {name: ToTensor} - - {name: Normalize, mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225]} - loader: - batch_size: 16 - shuffle: false - num_workers: 4 diff --git a/configs/eval/baseline.yaml b/configs/eval/baseline.yaml deleted file mode 100644 index fe22d55..0000000 --- a/configs/eval/baseline.yaml +++ /dev/null @@ -1,12 +0,0 @@ -_base_: - - ../train/baseline.yaml - -# auto 按 cuda > mps > cpu 自动选；部署前做精度回归可以显式 device=cpu。 -device: auto -checkpoint: runs/baseline/best.pth - -eval: - split: val - score_threshold: 0.3 - iou_threshold: 0.5 - topk: 50 diff --git a/configs/eval/resnet18.yaml b/configs/eval/resnet18.yaml deleted file mode 100644 index 3601819..0000000 --- a/configs/eval/resnet18.yaml +++ /dev/null @@ -1,13 +0,0 @@ -# resnet18 版本的评估配置。 -# 必须和训练时用的模型配置一致，否则构建出来的骨架与权重对不上。 -_base_: - - ../train/resnet18.yaml - -device: auto -checkpoint: runs/resnet18/best.pth - -eval: - split: val - score_threshold: 0.3 - iou_threshold: 0.5 - topk: 50 diff --git a/configs/model/resnet18.yaml b/configs/model/resnet18.yaml deleted file mode 100644 index f2df69d..0000000 --- a/configs/model/resnet18.yaml +++ /dev/null @@ -1,18 +0,0 @@ -# torchvision resnet18：stride 32 + 3 次上采样 => 输出 stride 4。 -model: - name: defect_net - backbone: - # 整块替换：不加这行的话，从 tiny_conv 继承来的 widths/depths/block 会残留， - # 传给 TorchvisionBackbone 直接报错。见 utils/config.py::deep_merge - _replace_: true - name: torchvision - arch: resnet18 - pretrained: true - in_channels: 3 - neck: - name: upsample - channels: 64 - num_upsamples: 3 - head: - hidden_channels: 64 - with_classification: true diff --git a/configs/model/tiny_conv.yaml b/configs/model/tiny_conv.yaml deleted file mode 100644 index 210c7ae..0000000 --- a/configs/model/tiny_conv.yaml +++ /dev/null @@ -1,18 +0,0 @@ -# 自研轻量骨架：backbone stride 16 + 2 次上采样 => 输出 stride 4。 -model: - name: defect_net - backbone: - name: tiny_conv - in_channels: 3 - widths: [32, 64, 128, 256] - depths: [1, 2, 2, 2] - block: dw # dw=深度可分离(更快) / res=残差(更强) - neck: - name: upsample - channels: 64 - num_upsamples: 2 - head: - hidden_channels: 64 - with_classification: true - classification: - dropout: 0.1 diff --git a/configs/train/baseline.yaml b/configs/train/baseline.yaml deleted file mode 100644 index d366433..0000000 --- a/configs/train/baseline.yaml +++ /dev/null @@ -1,43 +0,0 @@ -_base_: - - ../model/tiny_conv.yaml - - ../dataset/film_defect.yaml - -work_dir: runs/baseline -seed: 42 -deterministic: false -device: auto # auto 按 cuda > mps > cpu 挑最快的可用后端；也可写 cuda / mps / cpu - -train: - epochs: 100 - amp: auto # auto: CUDA 上开、MPS 上关（实测 MPS autocast 比 fp32 慢）；也可写 true / false - clip_grad: 5.0 - log_interval: 20 - monitor: f1 # 用于挑选 best.pth 的验证指标 - restart_schedule: false # 续训时按剩余轮数重开学习率曲线（调度已到末尾、想调大 lr 时用） - - optimizer: - name: adamw - lr: 0.001 - weight_decay: 0.0001 - - scheduler: - name: cosine - - loss: - heatmap_weight: 1.0 - wh_weight: 0.1 - offset_weight: 1.0 - cls_weight: 1.0 - - eval: - score_threshold: 0.3 - iou_threshold: 0.5 - topk: 50 - - # 训练过程可视化（SwanLab）。未安装或 enabled=false 时自动跳过。 - swanlab: - enabled: true - project: film-defect - name: null # 缺省用 work_dir 的目录名 - mode: online # online(云端) | local(本地看板) | offline(仅存盘) | disabled - tags: [tiny_conv, center-head] diff --git a/configs/train/multiclass.yaml b/configs/train/multiclass.yaml deleted file mode 100644 index 2bb1e3c..0000000 --- a/configs/train/multiclass.yaml +++ /dev/null @@ -1,12 +0,0 @@ -_base_: - - ./baseline.yaml - -# 单类（互斥）分类示例：一张图只归一个类别，分类头多一维表示"无缺陷"。 -label_mode: multiclass - -work_dir: runs/multiclass -train: - monitor: cls_f1 # 单类模式下常用分类指标挑选 best.pth - swanlab: - name: multiclass - tags: [tiny_conv, multiclass] diff --git a/configs/train/resnet18.yaml b/configs/train/resnet18.yaml deleted file mode 100644 index 1d79d02..0000000 --- a/configs/train/resnet18.yaml +++ /dev/null @@ -1,8 +0,0 @@ -_base_: - - ./baseline.yaml - - ../model/resnet18.yaml - -work_dir: runs/resnet18 -train: - optimizer: - lr: 0.0005 diff --git a/datasets/__init__.py b/datasets/__init__.py index 3f70885..1e301cd 100644 --- a/datasets/__init__.py +++ b/datasets/__init__.py @@ -1,28 +1,20 @@ -"""数据集：抛出样本的 Dataset 类与相应的数据变换、标签编码。""" +"""数据集：无标注（自监督/聚类用）与按目录组织的分类数据集。""" -from datasets.build import build_dataloader, build_dataset +from datasets.build import build_dataset from datasets.classification import FolderClassificationDataset, classification_collate -from datasets.film_defect import FilmDefectDataset, collate_fn from datasets.registry import DATASETS -from datasets.target_encoder import CenterTargetEncoder from datasets.transforms import build_transforms from datasets.unlabeled import UnlabeledImageDataset, unlabeled_collate -from datasets.yolo import format_label_lines, load_classes, parse_label_file, write_data_yaml +from datasets.yolo import IMAGE_SUFFIXES, write_data_yaml __all__ = [ "build_dataset", "FolderClassificationDataset", "classification_collate", - "build_dataloader", - "FilmDefectDataset", - "collate_fn", - "DATASETS", "UnlabeledImageDataset", "unlabeled_collate", - "CenterTargetEncoder", "build_transforms", - "parse_label_file", - "format_label_lines", - "load_classes", + "DATASETS", + "IMAGE_SUFFIXES", "write_data_yaml", ] diff --git a/datasets/build.py b/datasets/build.py index e950ba4..22925b1 100644 --- a/datasets/build.py +++ b/datasets/build.py @@ -1,39 +1,14 @@ -"""按配置构建 Dataset / DataLoader。""" +"""按配置构建 Dataset。""" from __future__ import annotations from collections.abc import Mapping from typing import Any -from torch.utils.data import DataLoader, Dataset +from torch.utils.data import Dataset -from datasets.film_defect import collate_fn from datasets.registry import DATASETS -from utils.device import pin_memory_supported, resolve_device def build_dataset(cfg: Mapping[str, Any], **kwargs: Any) -> Dataset: return DATASETS.build(dict(cfg), **kwargs) - - -def build_dataloader( - dataset: Dataset, - batch_size: int = 8, - shuffle: bool = False, - num_workers: int = 4, - pin_memory: bool | None = None, - drop_last: bool = False, -) -> DataLoader: - """`pin_memory` 缺省时只在 CUDA 上开启（CPU/MPS 上无收益且会告警）。""" - if pin_memory is None: - pin_memory = pin_memory_supported(resolve_device("auto")) - return DataLoader( - dataset, - batch_size=batch_size, - shuffle=shuffle, - num_workers=num_workers, - pin_memory=pin_memory, - drop_last=drop_last, - collate_fn=collate_fn, - persistent_workers=num_workers > 0, - ) diff --git a/datasets/film_defect.py b/datasets/film_defect.py deleted file mode 100644 index 483478c..0000000 --- a/datasets/film_defect.py +++ /dev/null @@ -1,131 +0,0 @@ -"""薄膜缺陷数据集（YOLO 风格标注）。 - -目录与标注格式见 `datasets/yolo.py` 与 data/README.md：`images/` 下的每张图， -在同层 `labels/` 下有一个同名 .txt，每行 `class_id cx cy w h`（相对宽高归一化）。 -缺失或空的 .txt 表示该图无缺陷。 - -图像级标签的形式由配置里的 `label_mode` 决定： - multilabel（多类）——多热向量，一张图可同时含多种缺陷，全 0 表示 OK 片； - multiclass（单类）——单个类别下标，互斥；无缺陷时取背景下标 len(classes)。 -用 scripts/convert_annotations.py 从 labelme / VOC 转成该格式。 -""" - -from __future__ import annotations - -from collections.abc import Mapping, Sequence -from pathlib import Path -from typing import Any - -import numpy as np -import torch -from PIL import Image -from torch.utils.data import Dataset - -from datasets.registry import DATASETS -from datasets.target_encoder import CenterTargetEncoder -from datasets.transforms import build_transforms -from datasets.yolo import denormalize, label_path_for, list_images, load_classes, parse_label_file -from utils.labels import MULTILABEL, check_label_mode - - -@DATASETS.register("film_defect") -class FilmDefectDataset(Dataset): - """Args: - img_dir: 图片目录（YOLO 约定，如 `data/raw/film/images/train`），递归查找图片。 - classes: 类别名列表，或指向 YOLO data.yaml 的路径（读其中的 `names`）。 - 列表顺序即类别下标顺序，必须与标注 txt 里的 class_id 对应。 - label_dir: 标注目录；缺省按 YOLO 约定把路径里的 `images` 换成 `labels`。 - list_file: 可选的图片清单（YOLO 的 train.txt 风格），给了就按它的顺序取图。 - transforms: transform 配置列表，见 `datasets.transforms.build_transforms`。 - stride: 模型输出下采样倍数，用于生成 heatmap 监督。 - gray: 是否按单通道读取（薄膜线扫相机通常为灰度）。 - max_objects: 单图最多编码的目标数。 - label_mode: "multilabel"(多类，多热向量) 或 "multiclass"(单类，互斥下标)。 - 必须与模型配置的 label_mode 一致，train.py / eval.py 会统一下发。 - """ - - def __init__( - self, - img_dir: str | Path, - classes: Sequence[str] | str | Path, - label_dir: str | Path | None = None, - list_file: str | Path | None = None, - transforms: Sequence[Mapping[str, Any]] = (), - stride: int = 4, - gray: bool = False, - max_objects: int = 64, - label_mode: str = "multilabel", - ) -> None: - self.label_mode = check_label_mode(label_mode) - self.img_dir = Path(img_dir) - self.label_dir = Path(label_dir) if label_dir is not None else None - self.classes = load_classes(classes) - self.class_to_idx = {name: i for i, name in enumerate(self.classes)} - self.transforms = build_transforms(transforms) - self.gray = gray - self.encoder = CenterTargetEncoder( - num_classes=len(self.classes), stride=stride, max_objects=max_objects - ) - - self.images = list_images(self.img_dir, Path(list_file) if list_file else None) - if not self.images: - raise FileNotFoundError(f"{self.img_dir} 下没有找到图片") - - def label_file(self, image_path: Path) -> Path: - return label_path_for(image_path, self.img_dir, self.label_dir) - - def __len__(self) -> int: - return len(self.images) - - def __getitem__(self, index: int) -> dict: - path = self.images[index] - image = Image.open(path) - image = image.convert("L" if self.gray else "RGB") - width, height = image.size - - norm_boxes, labels = parse_label_file(self.label_file(path), num_classes=len(self.classes)) - boxes = denormalize(norm_boxes, width, height) - - sample = self.transforms({"image": image, "boxes": boxes, "labels": labels}) - image_tensor: torch.Tensor = sample["image"] - boxes = np.asarray(sample["boxes"], dtype=np.float32).reshape(-1, 4) - labels = np.asarray(sample["labels"], dtype=np.int64) - - target = self.encoder(boxes, labels, image_tensor.shape[-2:]) - image_label = self._image_label(labels) - - return { - "image": image_tensor, - "target": target, - "boxes": torch.from_numpy(boxes), - "labels": torch.from_numpy(labels), - "image_label": image_label, - "meta": {"path": str(path), "index": index}, - } - - def _image_label(self, labels: np.ndarray) -> torch.Tensor: - """多类 -> 多热向量 (num_classes,)；单类 -> 标量下标（背景为 len(classes)）。""" - if self.label_mode == MULTILABEL: - image_label = torch.zeros(len(self.classes), dtype=torch.float32) - if len(labels): - image_label[torch.from_numpy(labels)] = 1.0 - return image_label - - if not len(labels): - return torch.tensor(len(self.classes), dtype=torch.long) - # 一图多种缺陷时取框数最多的类别（并列取下标最小的），保证结果确定。 - counts = np.bincount(labels, minlength=len(self.classes)) - return torch.tensor(int(counts.argmax()), dtype=torch.long) - - -def collate_fn(batch: list[dict]) -> dict: - """框数量不定，按 list 保留；其余堆叠成 batch 张量。""" - target_keys = batch[0]["target"].keys() - return { - "images": torch.stack([item["image"] for item in batch]), - "target": {key: torch.stack([item["target"][key] for item in batch]) for key in target_keys}, - "boxes": [item["boxes"] for item in batch], - "labels": [item["labels"] for item in batch], - "image_labels": torch.stack([item["image_label"] for item in batch]), - "meta": [item["meta"] for item in batch], - } diff --git a/datasets/target_encoder.py b/datasets/target_encoder.py deleted file mode 100644 index f5dad83..0000000 --- a/datasets/target_encoder.py +++ /dev/null @@ -1,111 +0,0 @@ -"""把框标注编码成 CenterHead 需要的监督信号。""" - -from __future__ import annotations - -import numpy as np -import torch - - -def gaussian_radius(height: float, width: float, min_overlap: float = 0.7) -> float: - """CenterNet 的半径公式：保证半径内的框与 GT 的 IoU >= min_overlap。""" - a1 = 1.0 - b1 = height + width - c1 = width * height * (1 - min_overlap) / (1 + min_overlap) - r1 = (b1 - np.sqrt(max(b1**2 - 4 * a1 * c1, 0.0))) / (2 * a1) - - a2 = 4.0 - b2 = 2 * (height + width) - c2 = (1 - min_overlap) * width * height - r2 = (b2 - np.sqrt(max(b2**2 - 4 * a2 * c2, 0.0))) / (2 * a2) - - a3 = 4 * min_overlap - b3 = -2 * min_overlap * (height + width) - c3 = (min_overlap - 1) * width * height - r3 = (-b3 + np.sqrt(max(b3**2 - 4 * a3 * c3, 0.0))) / (2 * a3) - - return max(0.0, min(r1, r2, r3)) - - -def draw_gaussian(heatmap: np.ndarray, center: tuple[int, int], radius: int) -> None: - """在 heatmap 上叠加二维高斯（取 max，避免相邻目标互相削弱）。""" - diameter = 2 * radius + 1 - sigma = diameter / 6.0 - grid = np.arange(-radius, radius + 1, dtype=np.float32) - gaussian = np.exp(-(grid[None, :] ** 2 + grid[:, None] ** 2) / (2 * sigma**2)) - - cx, cy = center - height, width = heatmap.shape - left, right = min(cx, radius), min(width - cx, radius + 1) - top, bottom = min(cy, radius), min(height - cy, radius + 1) - if right <= -left or bottom <= -top: - return - - masked_heatmap = heatmap[cy - top : cy + bottom, cx - left : cx + right] - masked_gaussian = gaussian[radius - top : radius + bottom, radius - left : radius + right] - np.maximum(masked_heatmap, masked_gaussian, out=masked_heatmap) - - -class CenterTargetEncoder: - """生成 heatmap / wh / offset / mask 四个监督张量。 - - Args: - num_classes: 类别数。 - stride: 模型输出相对输入的下采样倍数。 - max_objects: 单图最多编码的目标数。 - min_overlap: 高斯半径的 IoU 阈值。 - """ - - def __init__( - self, - num_classes: int, - stride: int, - max_objects: int = 64, - min_overlap: float = 0.7, - ) -> None: - self.num_classes = num_classes - self.stride = stride - self.max_objects = max_objects - self.min_overlap = min_overlap - - def __call__( - self, - boxes: np.ndarray, - labels: np.ndarray, - image_size: tuple[int, int], - ) -> dict[str, torch.Tensor]: - img_h, img_w = image_size - out_h, out_w = img_h // self.stride, img_w // self.stride - - heatmap = np.zeros((self.num_classes, out_h, out_w), dtype=np.float32) - wh = np.zeros((2, out_h, out_w), dtype=np.float32) - offset = np.zeros((2, out_h, out_w), dtype=np.float32) - mask = np.zeros((out_h, out_w), dtype=np.float32) - - count = 0 - for box, label in zip(np.asarray(boxes, dtype=np.float32), np.asarray(labels), strict=True): - if count >= self.max_objects: - break - x1, y1, x2, y2 = box / self.stride - box_w, box_h = x2 - x1, y2 - y1 - if box_w <= 0 or box_h <= 0: - continue - - cx, cy = (x1 + x2) / 2, (y1 + y2) / 2 - cx_int, cy_int = int(cx), int(cy) - if not (0 <= cx_int < out_w and 0 <= cy_int < out_h): - continue - - radius = max(1, int(gaussian_radius(box_h, box_w, self.min_overlap))) - draw_gaussian(heatmap[int(label)], (cx_int, cy_int), radius) - - wh[:, cy_int, cx_int] = (box_w, box_h) - offset[:, cy_int, cx_int] = (cx - cx_int, cy - cy_int) - mask[cy_int, cx_int] = 1.0 - count += 1 - - return { - "heatmap": torch.from_numpy(heatmap), - "wh": torch.from_numpy(wh), - "offset": torch.from_numpy(offset), - "mask": torch.from_numpy(mask), - } diff --git a/datasets/yolo.py b/datasets/yolo.py index 4b0adee..6a06e95 100644 --- a/datasets/yolo.py +++ b/datasets/yolo.py @@ -1,24 +1,13 @@ -"""YOLO 风格标注的读写。 +"""YOLO 数据集格式的共用定义。 -目录约定（与 ultralytics 一致，`images` 与 `labels` 同层对应）： +工作流里只用到"写出"这一侧：scripts/generate_boxes.py 产出的 +`labels/<类别>/<图名>.txt` 与 `data.yaml` 是给外部检测训练用的交换格式。 -``` -data/raw/film/ -├── images/train/0001.png -├── labels/train/0001.txt # 与图片同名，后缀换成 .txt -├── images/val/... -├── labels/val/... -└── data.yaml # path / train / val / names -``` - -标注文件每行一个目标，坐标是**相对图像宽高归一化**到 [0, 1] 的中心点格式： +标注文件每行一个目标，坐标相对图像宽高归一化到 [0, 1]，中心点格式： ``` ``` - -也接受 YOLO 分割格式的多边形（` x1 y1 x2 y2 ...`，偶数个坐标）， -读取时取其外接矩形。空文件或缺失文件表示该图无缺陷（纯负样本）。 """ from __future__ import annotations @@ -26,140 +15,9 @@ from __future__ import annotations from collections.abc import Sequence from pathlib import Path -import numpy as np - IMAGE_SUFFIXES = (".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp") -def label_path_for(image_path: Path, img_dir: Path, label_dir: Path | None = None) -> Path: - """图片路径 -> 标注路径。 - - 缺省按 YOLO 约定把路径里最后一个 `images` 目录名换成 `labels`； - 显式给了 `label_dir` 时，保持图片相对 `img_dir` 的子目录结构。 - """ - relative = image_path.relative_to(img_dir).with_suffix(".txt") - if label_dir is not None: - return Path(label_dir) / relative - - parts = list(img_dir.parts) - for i in range(len(parts) - 1, -1, -1): - if parts[i] == "images": - parts[i] = "labels" - return Path(*parts) / relative - # 没有 images 目录时退化为同级的 labels/ - return img_dir.parent / "labels" / relative - - -def list_images(img_dir: Path, list_file: Path | None = None) -> list[Path]: - """列出图片；给了 `list_file` 则按文件里的路径顺序（YOLO 的 train.txt 风格）。""" - if list_file is not None: - list_file = Path(list_file) - if not list_file.exists(): - raise FileNotFoundError(f"图片清单不存在: {list_file}") - images = [] - for line in list_file.read_text(encoding="utf-8").splitlines(): - line = line.strip() - if not line or line.startswith("#"): - continue - path = Path(line) - images.append(path if path.is_absolute() else (list_file.parent / path).resolve()) - return images - - if not img_dir.is_dir(): - raise FileNotFoundError(f"图片目录不存在: {img_dir}") - return sorted(p for p in img_dir.rglob("*") if p.suffix.lower() in IMAGE_SUFFIXES) - - -def parse_label_file(path: Path, num_classes: int | None = None) -> tuple[np.ndarray, np.ndarray]: - """读取一个 .txt，返回归一化的 xyxy 框 (N, 4) 与类别下标 (N,)。 - - 缺失文件按无缺陷处理（YOLO 的惯例），格式错误则带上 `文件:行号` 抛出。 - """ - boxes: list[list[float]] = [] - labels: list[int] = [] - if not path.exists(): - return np.zeros((0, 4), dtype=np.float32), np.zeros((0,), dtype=np.int64) - - for line_no, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1): - line = line.strip() - if not line or line.startswith("#"): - continue - parts = line.split() - try: - class_id = int(float(parts[0])) - values = [float(v) for v in parts[1:]] - except ValueError as exc: - raise ValueError(f"{path}:{line_no} 不是合法的 YOLO 标注行: {line!r}") from exc - - if num_classes is not None and not 0 <= class_id < num_classes: - raise ValueError( - f"{path}:{line_no} 类别下标 {class_id} 越界，classes 只有 {num_classes} 个" - ) - - if len(values) == 4: - cx, cy, width, height = values - box = [cx - width / 2, cy - height / 2, cx + width / 2, cy + height / 2] - elif len(values) >= 6 and len(values) % 2 == 0: - # 分割格式的多边形，取外接矩形 - xs, ys = values[0::2], values[1::2] - box = [min(xs), min(ys), max(xs), max(ys)] - else: - raise ValueError( - f"{path}:{line_no} 期望 `class cx cy w h` 或分割多边形，实际有 {len(values)} 个坐标" - ) - - box = [min(max(v, 0.0), 1.0) for v in box] - if box[2] <= box[0] or box[3] <= box[1]: - continue # 退化框直接丢弃 - boxes.append(box) - labels.append(class_id) - - return ( - np.asarray(boxes, dtype=np.float32).reshape(-1, 4), - np.asarray(labels, dtype=np.int64), - ) - - -def denormalize(boxes: np.ndarray, width: int, height: int) -> np.ndarray: - """归一化 xyxy -> 像素 xyxy。""" - if not len(boxes): - return np.zeros((0, 4), dtype=np.float32) - scale = np.array([width, height, width, height], dtype=np.float32) - return boxes.astype(np.float32) * scale - - -def format_label_lines(boxes: np.ndarray, labels: Sequence[int], width: int, height: int) -> list[str]: - """像素 xyxy -> YOLO 文本行（归一化 cx cy w h，保留 6 位小数）。""" - lines = [] - for (x1, y1, x2, y2), label in zip(np.asarray(boxes, dtype=np.float64), labels, strict=True): - x1, x2 = sorted((max(x1, 0.0), min(x2, float(width)))) - y1, y2 = sorted((max(y1, 0.0), min(y2, float(height)))) - box_w, box_h = (x2 - x1) / width, (y2 - y1) / height - if box_w <= 0 or box_h <= 0: - continue - cx, cy = (x1 + x2) / 2 / width, (y1 + y2) / 2 / height - lines.append(f"{int(label)} {cx:.6f} {cy:.6f} {box_w:.6f} {box_h:.6f}") - return lines - - -def load_classes(spec: Sequence[str] | str | Path) -> list[str]: - """类别名可以直接写成列表，也可以指向 YOLO 的 data.yaml（读其中的 `names`）。""" - if not isinstance(spec, str | Path): - return list(spec) - - import yaml - - path = Path(spec) - if not path.exists(): - raise FileNotFoundError(f"data.yaml 不存在: {path}") - names = (yaml.safe_load(path.read_text(encoding="utf-8")) or {}).get("names") - if names is None: - raise KeyError(f"{path} 里没有 `names` 字段") - if isinstance(names, dict): # {0: scratch, 1: bubble} - return [names[key] for key in sorted(names, key=int)] - return list(names) - - def write_data_yaml( path: str | Path, classes: Sequence[str], @@ -167,7 +25,7 @@ def write_data_yaml( train: str = "images/train", val: str = "images/val", ) -> Path: - """写一个 YOLO 风格的 data.yaml。""" + """写一个 YOLO 风格的 data.yaml，字段与 ultralytics 通用。""" import yaml path = Path(path) @@ -179,7 +37,7 @@ def write_data_yaml( "train": train, "val": val, "nc": len(classes), - "names": {i: name for i, name in enumerate(classes)}, + "names": dict(enumerate(classes)), }, allow_unicode=True, sort_keys=False, diff --git a/engine/__init__.py b/engine/__init__.py index 2c5e4f7..22b72e1 100644 --- a/engine/__init__.py +++ b/engine/__init__.py @@ -1,4 +1,4 @@ -"""训练 / 评估流程。""" +"""训练与分析流程：自监督、聚类、分类、归因定位。""" from engine.checkpoint import load_checkpoint, save_checkpoint from engine.classification import ( @@ -6,18 +6,14 @@ from engine.classification import ( ClassificationTrainer, class_weights, ) -from engine.evaluator import Evaluator -from engine.losses import DetectionLoss from engine.optim import build_optimizer, build_scheduler -from engine.trainer import Trainer +from engine.simclr import SimCLRTrainer __all__ = [ - "Trainer", "ClassificationTrainer", "ClassificationMetric", "class_weights", - "Evaluator", - "DetectionLoss", + "SimCLRTrainer", "build_optimizer", "build_scheduler", "save_checkpoint", diff --git a/engine/evaluator.py b/engine/evaluator.py deleted file mode 100644 index f3e8a4f..0000000 --- a/engine/evaluator.py +++ /dev/null @@ -1,87 +0,0 @@ -"""评估：跑一遍验证集，输出分类与定位指标。""" - -from __future__ import annotations - -from typing import Any - -import torch -from torch import nn -from torch.utils.data import DataLoader - -from engine.metrics import AverageMeter, DetectionMetric, MultiLabelMetric, SingleLabelMetric -from models.postprocess import decode_heatmap -from utils.labels import MULTILABEL -from utils.logging import get_logger - -logger = get_logger("engine.evaluator") - - -class Evaluator: - def __init__( - self, - model: nn.Module, - dataloader: DataLoader, - device: torch.device, - criterion: nn.Module | None = None, - score_threshold: float = 0.3, - iou_threshold: float = 0.5, - topk: int = 50, - ) -> None: - self.model = model - self.dataloader = dataloader - self.device = device - self.criterion = criterion - self.score_threshold = score_threshold - self.topk = topk - self.det_metric = DetectionMetric(iou_threshold) - # 分类指标按模型的 label_mode 选择：多类用多标签指标，单类用互斥分类指标。 - self.label_mode = getattr(model, "label_mode", MULTILABEL) - if self.label_mode == MULTILABEL: - self.cls_metric: MultiLabelMetric | SingleLabelMetric = MultiLabelMetric() - else: - num_outputs = model.cls_head.num_outputs if model.cls_head is not None else 1 - self.cls_metric = SingleLabelMetric(num_outputs) - - @torch.no_grad() - def evaluate(self) -> dict[str, float]: - self.model.eval() - self.det_metric.reset() - self.cls_metric.reset() - loss_meter = AverageMeter() - - stride = getattr(self.model, "out_stride", 4) - for batch in self.dataloader: - images = batch["images"].to(self.device, non_blocking=True) - outputs = self.model(images) - - if self.criterion is not None: - batch_on_device = _to_device(batch, self.device) - losses = self.criterion(outputs, batch_on_device) - loss_meter.update(float(losses["total"]), images.shape[0]) - - detections = decode_heatmap( - outputs["heatmap"], - outputs["wh"], - outputs["offset"], - stride=stride, - topk=self.topk, - score_threshold=self.score_threshold, - ) - for i, det in enumerate(detections): - self.det_metric.update(det["boxes"], det["labels"], batch["boxes"][i], batch["labels"][i]) - - if "logits" in outputs: - self.cls_metric.update(outputs["logits"], batch["image_labels"]) - - metrics: dict[str, float] = {**self.det_metric.compute(), **self.cls_metric.compute()} - if self.criterion is not None: - metrics["loss"] = loss_meter.avg - return metrics - - -def _to_device(batch: dict, device: torch.device) -> dict: - moved: dict[str, Any] = dict(batch) - moved["images"] = batch["images"].to(device, non_blocking=True) - moved["target"] = {k: v.to(device, non_blocking=True) for k, v in batch["target"].items()} - moved["image_labels"] = batch["image_labels"].to(device, non_blocking=True) - return moved diff --git a/engine/losses.py b/engine/losses.py deleted file mode 100644 index a58b8a3..0000000 --- a/engine/losses.py +++ /dev/null @@ -1,79 +0,0 @@ -"""损失函数。""" - -from __future__ import annotations - -import torch -from torch import nn -from torch.nn import functional as F - -from utils.labels import MULTILABEL, check_label_mode - - -def gaussian_focal_loss( - pred_logits: torch.Tensor, - gt_heatmap: torch.Tensor, - alpha: float = 2.0, - beta: float = 4.0, - eps: float = 1e-4, -) -> torch.Tensor: - """CenterNet 的 penalty-reduced focal loss，按正样本数归一化。""" - pred = pred_logits.sigmoid().clamp(eps, 1 - eps) - pos_mask = gt_heatmap.eq(1.0).float() - neg_mask = 1.0 - pos_mask - - pos_loss = -((1 - pred) ** alpha) * torch.log(pred) * pos_mask - neg_loss = -((1 - gt_heatmap) ** beta) * (pred**alpha) * torch.log(1 - pred) * neg_mask - - num_pos = pos_mask.sum().clamp(min=1.0) - return (pos_loss.sum() + neg_loss.sum()) / num_pos - - -def masked_l1_loss(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: - """只在中心点位置计算 L1。mask: (B, H, W)。""" - mask = mask.unsqueeze(1).expand_as(pred) - loss = F.l1_loss(pred * mask, target * mask, reduction="sum") - return loss / mask.sum().clamp(min=1.0) - - -class DetectionLoss(nn.Module): - """heatmap + wh + offset (+ 图像级分类) 的加权和。 - - 分类分支按 `label_mode` 切换： - multilabel（多类）—— 多热标签，逐类 BCE； - multiclass（单类）—— 互斥下标，softmax 交叉熵。 - """ - - def __init__( - self, - heatmap_weight: float = 1.0, - wh_weight: float = 0.1, - offset_weight: float = 1.0, - cls_weight: float = 1.0, - label_mode: str = MULTILABEL, - ) -> None: - super().__init__() - self.label_mode = check_label_mode(label_mode) - self.heatmap_weight = heatmap_weight - self.wh_weight = wh_weight - self.offset_weight = offset_weight - self.cls_weight = cls_weight - - def forward(self, outputs: dict, batch: dict) -> dict[str, torch.Tensor]: - target = batch["target"] - losses = { - "heatmap": self.heatmap_weight * gaussian_focal_loss(outputs["heatmap"], target["heatmap"]), - "wh": self.wh_weight * masked_l1_loss(outputs["wh"], target["wh"], target["mask"]), - "offset": self.offset_weight - * masked_l1_loss(outputs["offset"], target["offset"], target["mask"]), - } - if "logits" in outputs and self.cls_weight > 0: - losses["cls"] = self.cls_weight * self._classification_loss( - outputs["logits"], batch["image_labels"] - ) - losses["total"] = sum(losses.values()) - return losses - - def _classification_loss(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: - if self.label_mode == MULTILABEL: - return F.binary_cross_entropy_with_logits(logits, targets.float()) - return F.cross_entropy(logits, targets.long()) diff --git a/engine/metrics.py b/engine/metrics.py index 56d3b9f..4a9a563 100644 --- a/engine/metrics.py +++ b/engine/metrics.py @@ -1,9 +1,10 @@ -"""指标：分类与定位。""" +"""训练过程中的通用统计。 + +分类任务的指标（混淆矩阵、macro-F1 等）在 engine/classification.py 里。 +""" from __future__ import annotations -import torch - class AverageMeter: def __init__(self) -> None: @@ -17,155 +18,3 @@ class AverageMeter: @property def avg(self) -> float: return self.total / self.count if self.count else 0.0 - - -def box_iou(boxes_a: torch.Tensor, boxes_b: torch.Tensor) -> torch.Tensor: - """(N, 4) 与 (M, 4) 的两两 IoU，返回 (N, M)。""" - if boxes_a.numel() == 0 or boxes_b.numel() == 0: - return boxes_a.new_zeros((boxes_a.shape[0], boxes_b.shape[0])) - area_a = (boxes_a[:, 2] - boxes_a[:, 0]).clamp(min=0) * (boxes_a[:, 3] - boxes_a[:, 1]).clamp(min=0) - area_b = (boxes_b[:, 2] - boxes_b[:, 0]).clamp(min=0) * (boxes_b[:, 3] - boxes_b[:, 1]).clamp(min=0) - - lt = torch.max(boxes_a[:, None, :2], boxes_b[None, :, :2]) - rb = torch.min(boxes_a[:, None, 2:], boxes_b[None, :, 2:]) - wh = (rb - lt).clamp(min=0) - inter = wh[..., 0] * wh[..., 1] - return inter / (area_a[:, None] + area_b[None, :] - inter).clamp(min=1e-6) - - -class DetectionMetric: - """按 IoU 阈值做贪心匹配，统计 precision / recall / F1 与平均 IoU。 - - 这是工程上够用的近似指标；正式对比模型时建议再跑一次 COCO mAP。 - """ - - def __init__(self, iou_threshold: float = 0.5) -> None: - self.iou_threshold = iou_threshold - self.reset() - - def reset(self) -> None: - self.tp = 0 - self.fp = 0 - self.fn = 0 - self.iou_sum = 0.0 - - def update( - self, - pred_boxes: torch.Tensor, - pred_labels: torch.Tensor, - gt_boxes: torch.Tensor, - gt_labels: torch.Tensor, - ) -> None: - if gt_boxes.numel() == 0: - self.fp += int(pred_boxes.shape[0]) - return - if pred_boxes.numel() == 0: - self.fn += int(gt_boxes.shape[0]) - return - - ious = box_iou(pred_boxes.float().cpu(), gt_boxes.float().cpu()) - same_class = pred_labels.cpu()[:, None] == gt_labels.cpu()[None, :] - ious = ious * same_class.float() - - matched_gt: set[int] = set() - for pred_idx in range(ious.shape[0]): - best_iou, best_gt = ious[pred_idx].max(dim=0) - gt_idx = int(best_gt) - if float(best_iou) >= self.iou_threshold and gt_idx not in matched_gt: - matched_gt.add(gt_idx) - self.tp += 1 - self.iou_sum += float(best_iou) - else: - self.fp += 1 - self.fn += int(gt_boxes.shape[0]) - len(matched_gt) - - def compute(self) -> dict[str, float]: - precision = self.tp / max(self.tp + self.fp, 1) - recall = self.tp / max(self.tp + self.fn, 1) - f1 = 2 * precision * recall / max(precision + recall, 1e-9) - return { - "precision": precision, - "recall": recall, - "f1": f1, - "mean_iou": self.iou_sum / max(self.tp, 1), - } - - -class MultiLabelMetric: - """多类模式：逐类阈值化后统计 micro precision / recall / F1 与整图完全正确率。""" - - def __init__(self, threshold: float = 0.5) -> None: - self.threshold = threshold - self.reset() - - def reset(self) -> None: - self.tp = 0 - self.fp = 0 - self.fn = 0 - self.exact = 0 - self.total = 0 - - def update(self, logits: torch.Tensor, targets: torch.Tensor) -> None: - pred = (logits.sigmoid() >= self.threshold).float().cpu() - gt = targets.float().cpu() - self.tp += int((pred * gt).sum()) - self.fp += int((pred * (1 - gt)).sum()) - self.fn += int(((1 - pred) * gt).sum()) - self.exact += int((pred == gt).all(dim=1).sum()) - self.total += int(gt.shape[0]) - - def compute(self) -> dict[str, float]: - precision = self.tp / max(self.tp + self.fp, 1) - recall = self.tp / max(self.tp + self.fn, 1) - return { - "cls_precision": precision, - "cls_recall": recall, - "cls_f1": 2 * precision * recall / max(precision + recall, 1e-9), - "cls_exact_match": self.exact / max(self.total, 1), - } - - -class SingleLabelMetric: - """单类模式：互斥分类的准确率与 macro precision / recall / F1。 - - 指标名与 `MultiLabelMetric` 保持一致（cls_precision / cls_recall / cls_f1）， - 因此切换 label_mode 后 `train.monitor` 不需要改。 - """ - - def __init__(self, num_outputs: int) -> None: - self.num_outputs = num_outputs - self.reset() - - def reset(self) -> None: - self.tp = torch.zeros(self.num_outputs) - self.fp = torch.zeros(self.num_outputs) - self.fn = torch.zeros(self.num_outputs) - self.correct = 0 - self.total = 0 - - def update(self, logits: torch.Tensor, targets: torch.Tensor) -> None: - pred = logits.argmax(dim=1).cpu() - gt = targets.long().cpu() - for cls in range(self.num_outputs): - pred_cls = pred == cls - gt_cls = gt == cls - self.tp[cls] += int((pred_cls & gt_cls).sum()) - self.fp[cls] += int((pred_cls & ~gt_cls).sum()) - self.fn[cls] += int((~pred_cls & gt_cls).sum()) - self.correct += int((pred == gt).sum()) - self.total += int(gt.numel()) - - def compute(self) -> dict[str, float]: - # 只对验证集中出现过的类别做 macro 平均，避免缺席类别把指标拉到 0。 - present = (self.tp + self.fn) > 0 - precision = self.tp / (self.tp + self.fp).clamp(min=1) - recall = self.tp / (self.tp + self.fn).clamp(min=1) - f1 = 2 * precision * recall / (precision + recall).clamp(min=1e-9) - if present.any(): - precision, recall, f1 = precision[present], recall[present], f1[present] - return { - "cls_accuracy": self.correct / max(self.total, 1), - "cls_precision": float(precision.mean()), - "cls_recall": float(recall.mean()), - "cls_f1": float(f1.mean()), - } diff --git a/engine/trainer.py b/engine/trainer.py deleted file mode 100644 index 4f27661..0000000 --- a/engine/trainer.py +++ /dev/null @@ -1,303 +0,0 @@ -"""训练循环。支持 cuda / mps / cpu 三种后端。""" - -from __future__ import annotations - -import math -import time -import warnings -from collections.abc import Mapping -from pathlib import Path -from typing import Any - -import torch -from torch import nn -from torch.utils.data import DataLoader - -from engine.checkpoint import load_checkpoint, save_checkpoint -from engine.evaluator import Evaluator, _to_device -from engine.losses import DetectionLoss -from engine.metrics import AverageMeter -from engine.optim import build_optimizer, build_scheduler -from engine.visualizer import SwanLabLogger -from utils.device import PeakMemory, resolve_amp -from utils.labels import MULTILABEL -from utils.logging import get_logger - -logger = get_logger("engine.trainer") - - -class Trainer: - """Args: - model: 已经 `.to(device)` 的模型。 - train_loader / val_loader: 数据加载器，val 可为 None。 - cfg: `train` 段配置（epochs / optimizer / scheduler / amp / ...）。 - work_dir: 日志与权重输出目录。 - """ - - def __init__( - self, - model: nn.Module, - train_loader: DataLoader, - val_loader: DataLoader | None, - cfg: Mapping[str, Any], - device: torch.device, - work_dir: str | Path, - full_config: dict | None = None, - ) -> None: - self.model = model - self.train_loader = train_loader - self.val_loader = val_loader - self.cfg = dict(cfg) - self.device = device - self.work_dir = Path(work_dir) - self.work_dir.mkdir(parents=True, exist_ok=True) - self.full_config = full_config - - self.epochs = int(self.cfg.get("epochs", 50)) - self.log_interval = int(self.cfg.get("log_interval", 20)) - self.clip_grad = float(self.cfg.get("clip_grad", 0.0)) - self.monitor = str(self.cfg.get("monitor", "f1")) - - # 分类损失形式跟随模型的 label_mode（yaml 里的 label_mode 决定多类 / 单类）。 - loss_cfg = dict(self.cfg.get("loss", {})) - loss_cfg.setdefault("label_mode", getattr(model, "label_mode", MULTILABEL)) - self.criterion = DetectionLoss(**loss_cfg) - self.optimizer = build_optimizer(model, self.cfg.get("optimizer", {"name": "adamw", "lr": 1e-3})) - self.scheduler = build_scheduler(self.optimizer, self.cfg.get("scheduler"), self.epochs) - - # amp: auto 时随后端（CUDA 开、MPS 关——MPS 上 autocast 实测比 fp32 慢） - self.amp = resolve_amp(self.cfg.get("amp", "auto"), device) - self.scaler = torch.amp.GradScaler(device.type, enabled=self.amp) - self.peak_memory = PeakMemory(device) - - self.evaluator = ( - Evaluator( - model, - val_loader, - device, - criterion=self.criterion, - **dict(self.cfg.get("eval", {})), - ) - if val_loader is not None - else None - ) - self.visualizer = SwanLabLogger( - work_dir=self.work_dir, - config=full_config, - **dict(self.cfg.get("swanlab", {})), - ) - self.best_metric = float("-inf") - self.start_epoch = 0 - self.global_step = 0 - - def train_one_epoch(self, epoch: int) -> dict[str, float]: - self.model.train() - meters: dict[str, AverageMeter] = {} - start = time.time() - - for step, batch in enumerate(self.train_loader): - batch = _to_device(batch, self.device) - - with torch.amp.autocast(self.device.type, enabled=self.amp): - outputs = self.model(batch["images"]) - losses = self.criterion(outputs, batch) - - self.optimizer.zero_grad(set_to_none=True) - self.scaler.scale(losses["total"]).backward() - if self.clip_grad > 0: - self.scaler.unscale_(self.optimizer) - nn.utils.clip_grad_norm_(self.model.parameters(), self.clip_grad) - self.scaler.step(self.optimizer) - self.scaler.update() - - batch_size = batch["images"].shape[0] - step_losses = {key: float(value.detach()) for key, value in losses.items()} - for key, value in step_losses.items(): - meters.setdefault(key, AverageMeter()).update(value, batch_size) - - self.global_step += 1 - if step % self.log_interval == 0: - self.visualizer.log( - {**step_losses, "lr": self.optimizer.param_groups[0]["lr"]}, - step=self.global_step, - prefix="train/", - ) - summary = " ".join(f"{k}={m.avg:.4f}" for k, m in meters.items()) - logger.info( - "epoch %d [%d/%d] lr=%.2e %s", - epoch, - step, - len(self.train_loader), - self.optimizer.param_groups[0]["lr"], - summary, - ) - - logger.info("epoch %d 训练耗时 %.1fs", epoch, time.time() - start) - return {key: meter.avg for key, meter in meters.items()} - - def _gpu_memory(self) -> dict[str, float]: - """本轮训练的显存峰值（GB）；推理端的详细分析见 scripts/profile_inference.py。 - - CUDA 读内置峰值计数器；MPS 没有峰值 API，用逐 step 采样的最大值（是下界）； - CPU 不统计。 - """ - if self.device.type == "cpu": - return {} - stats = { - "gpu_mem_allocated_gb": self.peak_memory.peak_mb / 1024, - "gpu_mem_reserved_gb": self.peak_memory.reserved_mb / 1024, - } - self.peak_memory.reset() - return stats - - def resume(self, path: str | Path) -> dict: - """从检查点续训。 - - **配置文件是唯一真相源**：检查点只提供"进度"——模型权重、优化器动量缓冲、 - 已训到第几个 epoch、global_step、best_metric；而 lr、weight_decay、总轮数、 - 调度器一律以当前 yaml 为准。否则改了配置再续训，改动会被检查点静默覆盖。 - """ - state = load_checkpoint(path, self.model, map_location=self.device) - - if state.get("optimizer"): - # optimizer.load_state_dict 会整体覆盖 param_groups，把 lr/weight_decay - # 一并还原成存档时的值。先抄下配置构造出的超参，加载完按组写回， - # 只留下动量缓冲。按组写回也保住了 BN 与 bias 那组 weight_decay=0 的结构。 - intended = [ - {key: value for key, value in group.items() if key != "params"} - for group in self.optimizer.param_groups - ] - self.optimizer.load_state_dict(state["optimizer"]) - for group, wanted in zip(self.optimizer.param_groups, intended, strict=True): - group.update(wanted) - - self.start_epoch = int(state.get("epoch", -1)) + 1 - self.global_step = int(state.get("global_step", 0)) - best = state.get("best_metric") - self.best_metric = float(best) if best is not None and math.isfinite(float(best)) else float("-inf") - - self._restore_scheduler() - self._log_config_changes(state) - logger.info( - "已从 %s 恢复：起始 epoch=%d，best_%s=%.4f", - path, - self.start_epoch, - self.monitor, - self.best_metric, - ) - return state - - def _restore_scheduler(self) -> None: - """按**当前配置**重建学习率调度，再快进到已训到的 epoch。 - - 不恢复检查点里的调度器状态：它带着当时的 T_max 与 base_lrs，一旦 yaml 里改了 - 总轮数或学习率就会被静默覆盖。cosine / step / multistep 都是 epoch 的确定性函数， - 快进与恢复状态等价，所以统一走重建这一条路。 - """ - if self.scheduler is None: - return - # cosine 的链式实现基于当前 lr 递推，必须先复位回 initial_lr， - # 否则从检查点带来的 lr 出发会算出一条错的曲线。 - for group in self.optimizer.param_groups: - group["lr"] = group.get("initial_lr", group["lr"]) - - if self.cfg.get("restart_schedule"): - # 续训到调度末段时 cosine 会把新设的 lr 乘上一个接近 0 的系数， - # "调大学习率接着练"就落空了。打开这个开关按剩余轮数重开一条曲线。 - remaining = max(1, self.epochs - self.start_epoch) - self.scheduler = build_scheduler(self.optimizer, self.cfg.get("scheduler"), remaining) - logger.warning( - "restart_schedule=true：学习率调度从 epoch %d 重新开始，覆盖剩余 %d 轮，起始 lr=%.2e", - self.start_epoch, - remaining, - self.optimizer.param_groups[0]["lr"], - ) - return - self.scheduler = build_scheduler(self.optimizer, self.cfg.get("scheduler"), self.epochs) - with warnings.catch_warnings(): - warnings.simplefilter("ignore") # 快进阶段的 "step before optimizer.step()" 属预期 - for _ in range(self.start_epoch): - self.scheduler.step() - - def _log_config_changes(self, state: dict) -> None: - """检查点里的配置与当前配置不一致时说清楚用的是哪一份。""" - archived = state.get("train_cfg") or {} - changes = [] - if archived.get("epochs") not in (None, self.epochs): - changes.append(f"总轮数 {archived['epochs']} -> {self.epochs}") - for key, value in (archived.get("optimizer") or {}).items(): - current = dict(self.cfg.get("optimizer", {})).get(key) - if current is not None and current != value: - changes.append(f"optimizer.{key} {value} -> {current}") - archived_scheduler = (archived.get("scheduler") or {}).get("name") - current_scheduler = (dict(self.cfg.get("scheduler") or {})).get("name") - if archived_scheduler and archived_scheduler != current_scheduler: - changes.append(f"scheduler {archived_scheduler} -> {current_scheduler}") - - if changes: - logger.warning("配置与检查点不一致，按**当前配置**执行：%s", "；".join(changes)) - logger.info( - "续训超参：lr=%.2e（initial_lr=%.2e）weight_decay=%s，共 %d 轮", - self.optimizer.param_groups[0]["lr"], - self.optimizer.param_groups[0].get("initial_lr", float("nan")), - [g["weight_decay"] for g in self.optimizer.param_groups], - self.epochs, - ) - - def fit(self) -> float: - logger.info("开始训练：共 %d 个 epoch，输出目录 %s", self.epochs, self.work_dir) - for epoch in range(self.start_epoch, self.epochs): - train_metrics = self.train_one_epoch(epoch) - if self.scheduler is not None: - self.scheduler.step() - - metrics: dict[str, float] = {} - if self.evaluator is not None: - metrics = self.evaluator.evaluate() - logger.info( - "epoch %d 验证: %s", - epoch, - " ".join(f"{k}={v:.4f}" for k, v in metrics.items()), - ) - - self.visualizer.log( - { - **{f"train/epoch_{k}": v for k, v in train_metrics.items()}, - **{f"val/{k}": v for k, v in metrics.items()}, - **{f"train/{k}": v for k, v in self._gpu_memory().items()}, - "epoch": epoch, - }, - step=self.global_step, - ) - - save_checkpoint( - self.work_dir / "last.pth", - self.model, - self.optimizer, - epoch=epoch, - best_metric=self.best_metric, - config=self.full_config, - scheduler=self.scheduler, - global_step=self.global_step, - train_cfg=dict(self.cfg), - ) - score = metrics.get(self.monitor) - if score is not None and score > self.best_metric: - self.best_metric = score - save_checkpoint( - self.work_dir / "best.pth", - self.model, - self.optimizer, - epoch=epoch, - best_metric=self.best_metric, - config=self.full_config, - scheduler=self.scheduler, - global_step=self.global_step, - train_cfg=dict(self.cfg), - ) - logger.info("epoch %d 刷新最佳 %s=%.4f，已保存 best.pth", epoch, self.monitor, score) - - if math.isfinite(self.best_metric): - self.visualizer.log({f"best_{self.monitor}": self.best_metric}, step=self.global_step) - self.visualizer.finish() - return self.best_metric diff --git a/eval.py b/eval.py deleted file mode 100755 index 57e7c2d..0000000 --- a/eval.py +++ /dev/null @@ -1,88 +0,0 @@ -#!/usr/bin/env python -"""评估入口。 - -用法: - python eval.py --config configs/eval/baseline.yaml --checkpoint runs/baseline/best.pth - -设备同样支持 auto / cuda / mps / cpu；部署前回归可以显式 `--device cpu`。 -""" - -from __future__ import annotations - -import argparse -import json -import sys -from pathlib import Path - -from datasets import build_dataloader, build_dataset, load_classes -from engine import Evaluator -from models import build_model -from utils import check_label_mode, load_config, setup_logging -from utils.device import DeviceNotAvailableError, describe_device, resolve_device - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="薄膜缺陷模型评估") - parser.add_argument("--config", required=True) - parser.add_argument("--checkpoint", default=None, help="覆盖配置中的 checkpoint") - parser.add_argument("--split", default=None, help="评估哪个 data 划分，默认取配置的 eval.split") - parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE") - parser.add_argument("--device", default=None, help="覆盖配置中的 device：auto / cuda / mps / cpu") - parser.add_argument("--out", default=None, help="把指标写入 json 文件") - return parser.parse_args() - - -def main() -> int: - args = parse_args() - cfg = load_config(args.config, args.set) - logger = setup_logging() - - eval_cfg = dict(cfg.get("eval", {})) - split = args.split or eval_cfg.pop("split", "val") - eval_cfg.pop("split", None) - - checkpoint = args.checkpoint or cfg.get("checkpoint") - if not checkpoint: - logger.error("必须通过 --checkpoint 或配置中的 checkpoint 指定权重") - return 1 - - try: - device = resolve_device(args.device or cfg.get("device", "auto")) - except DeviceNotAvailableError as exc: - logger.error("%s", exc) - return 1 - logger.info("评估设备: %s", describe_device(device)) - classes = load_classes(cfg["classes"]) - - # label_mode 由 yaml 决定，模型与数据集必须一致，这里统一下发。 - label_mode = check_label_mode(str(cfg.get("label_mode", "multilabel"))) - model_cfg = dict(cfg["model"]) - model_cfg.setdefault("num_classes", len(classes)) - model_cfg.setdefault("label_mode", label_mode) - model = build_model(model_cfg) - - from engine import load_checkpoint - - load_checkpoint(checkpoint, model, map_location=device) - model.to(device).eval() - - split_cfg = cfg["data"][split] - dataset = build_dataset( - split_cfg["dataset"], classes=classes, stride=model.out_stride, label_mode=label_mode - ) - loader = build_dataloader(dataset, **{**dict(split_cfg.get("loader", {})), "shuffle": False}) - logger.info("评估 %s 集，共 %d 张图，权重 %s", split, len(dataset), checkpoint) - - metrics = Evaluator(model, loader, device, **eval_cfg).evaluate() - for key, value in metrics.items(): - logger.info("%-16s %.4f", key, value) - - if args.out: - Path(args.out).parent.mkdir(parents=True, exist_ok=True) - Path(args.out).write_text(json.dumps(metrics, indent=2, ensure_ascii=False), encoding="utf-8") - logger.info("指标已写入 %s", args.out) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/models/__init__.py b/models/__init__.py index e540694..7fca8e6 100644 --- a/models/__init__.py +++ b/models/__init__.py @@ -1,22 +1,17 @@ -"""模型骨架：backbone -> neck -> heads。 +"""模型骨架：backbone + head。 -约束：部署在实时设备上，只使用基于 convolution 的结构（不含 attention / transformer）。 +按项目约定只用 convolution，不含 attention / transformer。 """ -from models import backbones, heads, necks # noqa: F401 触发注册 +from models import backbones, heads # noqa: F401 触发注册 from models.build import build_model from models.classifier import DefectClassifier -from models.defect_net import DefectNet -from models.postprocess import decode_heatmap -from models.registry import BACKBONES, HEADS, MODELS, NECKS +from models.registry import BACKBONES, HEADS, MODELS __all__ = [ "build_model", - "DefectNet", "DefectClassifier", - "decode_heatmap", "BACKBONES", - "NECKS", "HEADS", "MODELS", ] diff --git a/models/defect_net.py b/models/defect_net.py deleted file mode 100644 index 869d9f4..0000000 --- a/models/defect_net.py +++ /dev/null @@ -1,83 +0,0 @@ -"""缺陷分类 + 定位主模型。""" - -from __future__ import annotations - -from collections.abc import Mapping -from typing import Any - -import torch -from torch import nn - -from models.heads.center_head import CenterHead -from models.heads.classification import ClassificationHead -from models.registry import BACKBONES, MODELS, NECKS -from utils.labels import LABEL_MODES, MULTILABEL, check_label_mode - -__all__ = ["DefectNet", "classification_outputs", "LABEL_MODES"] - - -def classification_outputs(num_classes: int, label_mode: str) -> int: - """分类头的输出维度；multiclass 需要额外一维表示 "无缺陷"。""" - check_label_mode(label_mode) - return num_classes if label_mode == MULTILABEL else num_classes + 1 - - -@MODELS.register("defect_net") -class DefectNet(nn.Module): - """backbone -> neck -> {分类头, 中心点定位头}。 - - Args: - num_classes: 缺陷类别数（不含背景）。 - backbone: backbone 注册表的配置，如 `{"name": "tiny_conv", ...}`。 - neck: neck 配置，`None` 表示直接在骨架特征上出图。 - head: CenterHead 的额外参数。 - with_classification: 是否输出图像级分类 logits。 - label_mode: 图像级标签形式，由 yaml 配置决定。 - "multilabel" —— 多类：一张图可同时含多种缺陷，逐类 sigmoid + BCE； - "multiclass" —— 单类：一张图只归一个类别（含 "无缺陷"），softmax + CE。 - 仅影响分类头与分类损失；定位头始终按 num_classes 个类别输出 heatmap。 - """ - - def __init__( - self, - num_classes: int, - backbone: Mapping[str, Any] | None = None, - neck: Mapping[str, Any] | None = None, - head: Mapping[str, Any] | None = None, - with_classification: bool = True, - classification: Mapping[str, Any] | None = None, - label_mode: str = "multilabel", - ) -> None: - super().__init__() - self.num_classes = num_classes - self.label_mode = check_label_mode(label_mode) - - self.backbone = BACKBONES.build(dict(backbone or {"name": "tiny_conv"})) - stride = self.backbone.out_stride - channels = self.backbone.out_channels - - if neck: - self.neck: nn.Module = NECKS.build(dict(neck), in_channels=channels) - channels = self.neck.out_channels - stride = stride // self.neck.scale_factor - else: - self.neck = nn.Identity() - - self.out_stride = stride - self.center_head = CenterHead(channels, num_classes, **dict(head or {})) - self.cls_head = ( - ClassificationHead( - channels, - classification_outputs(num_classes, label_mode), - **dict(classification or {}), - ) - if with_classification - else None - ) - - def forward(self, images: torch.Tensor) -> dict[str, torch.Tensor]: - feat = self.neck(self.backbone(images)) - outputs = self.center_head(feat) - if self.cls_head is not None: - outputs["logits"] = self.cls_head(feat) - return outputs diff --git a/models/heads/__init__.py b/models/heads/__init__.py index c80cbec..c724ff5 100644 --- a/models/heads/__init__.py +++ b/models/heads/__init__.py @@ -1,4 +1,4 @@ -from models.heads.center_head import CenterHead from models.heads.classification import ClassificationHead +from models.heads.projection import ProjectionHead -__all__ = ["CenterHead", "ClassificationHead"] +__all__ = ["ClassificationHead", "ProjectionHead"] diff --git a/models/heads/center_head.py b/models/heads/center_head.py deleted file mode 100644 index 8c35613..0000000 --- a/models/heads/center_head.py +++ /dev/null @@ -1,51 +0,0 @@ -"""Anchor-free 中心点定位头（CenterNet 风格，纯卷积）。 - -输出: - heatmap: (B, num_classes, H, W) 各类别缺陷中心的置信度（logits） - wh: (B, 2, H, W) 中心点处的框宽高（特征图尺度） - offset: (B, 2, H, W) 下采样量化误差补偿 -""" - -from __future__ import annotations - -import math - -import torch -from torch import nn - -from models.layers import ConvBNAct -from models.registry import HEADS - - -def _branch(in_channels: int, hidden: int, out_channels: int) -> nn.Sequential: - return nn.Sequential( - ConvBNAct(in_channels, hidden, kernel_size=3), - nn.Conv2d(hidden, out_channels, kernel_size=1), - ) - - -@HEADS.register("center") -class CenterHead(nn.Module): - def __init__( - self, - in_channels: int, - num_classes: int, - hidden_channels: int = 64, - prior_prob: float = 0.01, - ) -> None: - super().__init__() - self.num_classes = num_classes - self.heatmap = _branch(in_channels, hidden_channels, num_classes) - self.wh = _branch(in_channels, hidden_channels, 2) - self.offset = _branch(in_channels, hidden_channels, 2) - - # heatmap 正样本极稀疏，用先验概率初始化最后一层 bias，稳定早期训练。 - bias = -math.log((1 - prior_prob) / prior_prob) - nn.init.constant_(self.heatmap[-1].bias, bias) - - def forward(self, feat: torch.Tensor) -> dict[str, torch.Tensor]: - return { - "heatmap": self.heatmap(feat), - "wh": self.wh(feat), - "offset": self.offset(feat), - } diff --git a/models/necks/__init__.py b/models/necks/__init__.py deleted file mode 100644 index f64a426..0000000 --- a/models/necks/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from models.necks.upsample import UpsampleNeck - -__all__ = ["UpsampleNeck"] diff --git a/models/necks/upsample.py b/models/necks/upsample.py deleted file mode 100644 index 3806bf6..0000000 --- a/models/necks/upsample.py +++ /dev/null @@ -1,33 +0,0 @@ -"""上采样 neck：把骨架的低分辨率特征恢复到定位所需的 stride。""" - -from __future__ import annotations - -import torch -from torch import nn -from torch.nn import functional as F - -from models.layers import ConvBNAct -from models.registry import NECKS - - -@NECKS.register("upsample") -class UpsampleNeck(nn.Module): - """`num_upsamples` 次 2 倍上采样，每次接一个 3x3 卷积平滑。 - - 使用 nearest + conv 而非转置卷积：避免棋盘伪影，且在推理引擎上更快。 - """ - - def __init__(self, in_channels: int, channels: int = 64, num_upsamples: int = 2) -> None: - super().__init__() - self.lateral = ConvBNAct(in_channels, channels, kernel_size=1) - self.blocks = nn.ModuleList( - [ConvBNAct(channels, channels, kernel_size=3) for _ in range(num_upsamples)] - ) - self.out_channels = channels - self.scale_factor = 2**num_upsamples - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.lateral(x) - for block in self.blocks: - x = block(F.interpolate(x, scale_factor=2.0, mode="nearest")) - return x diff --git a/models/postprocess.py b/models/postprocess.py deleted file mode 100644 index 91f3983..0000000 --- a/models/postprocess.py +++ /dev/null @@ -1,68 +0,0 @@ -"""推理后处理：从 heatmap 解码出缺陷框。""" - -from __future__ import annotations - -import torch -from torch.nn import functional as F - - -def heatmap_nms(heatmap: torch.Tensor, kernel: int = 3) -> torch.Tensor: - """用 max-pool 做峰值抑制，替代 NMS（无循环，可导出）。""" - pooled = F.max_pool2d(heatmap, kernel, stride=1, padding=kernel // 2) - return heatmap * (pooled == heatmap).float() - - -@torch.no_grad() -def decode_heatmap( - heatmap: torch.Tensor, - wh: torch.Tensor, - offset: torch.Tensor, - stride: int, - topk: int = 50, - score_threshold: float = 0.3, -) -> list[dict[str, torch.Tensor]]: - """解码为原图坐标下的检测结果。 - - Args: - heatmap: (B, C, H, W) logits。 - wh / offset: (B, 2, H, W)，单位为特征图像素。 - stride: 特征图相对原图的下采样倍数。 - - Returns: - 每张图一个 dict：`boxes` (N, 4) xyxy、`scores` (N,)、`labels` (N,)。 - """ - scores_map = heatmap_nms(heatmap.sigmoid()) - batch, _num_classes, height, width = scores_map.shape - - flat = scores_map.view(batch, -1) - topk = min(topk, flat.shape[1]) - topk_scores, topk_idx = flat.topk(topk, dim=1) - - labels = torch.div(topk_idx, height * width, rounding_mode="floor") - pixel_idx = topk_idx % (height * width) - ys = torch.div(pixel_idx, width, rounding_mode="floor").float() - xs = (pixel_idx % width).float() - - wh_flat = wh.view(batch, 2, -1) - offset_flat = offset.view(batch, 2, -1) - gather_idx = pixel_idx.unsqueeze(1).expand(-1, 2, -1) - pred_wh = wh_flat.gather(2, gather_idx) - pred_offset = offset_flat.gather(2, gather_idx) - - cx = (xs + pred_offset[:, 0]) * stride - cy = (ys + pred_offset[:, 1]) * stride - half_w = pred_wh[:, 0].clamp(min=0) * stride / 2 - half_h = pred_wh[:, 1].clamp(min=0) * stride / 2 - boxes = torch.stack([cx - half_w, cy - half_h, cx + half_w, cy + half_h], dim=-1) - - results = [] - for i in range(batch): - keep = topk_scores[i] >= score_threshold - results.append( - { - "boxes": boxes[i][keep], - "scores": topk_scores[i][keep], - "labels": labels[i][keep], - } - ) - return results diff --git a/models/registry.py b/models/registry.py index 6b968d0..77e0dc4 100644 --- a/models/registry.py +++ b/models/registry.py @@ -3,8 +3,7 @@ from utils.registry import Registry BACKBONES = Registry("backbones") -NECKS = Registry("necks") HEADS = Registry("heads") MODELS = Registry("models") -__all__ = ["BACKBONES", "NECKS", "HEADS", "MODELS"] +__all__ = ["BACKBONES", "HEADS", "MODELS"] diff --git a/scripts/convert_annotations.py b/scripts/convert_annotations.py deleted file mode 100755 index 6d65353..0000000 --- a/scripts/convert_annotations.py +++ /dev/null @@ -1,135 +0,0 @@ -#!/usr/bin/env python -"""把 labelme / VOC 标注批量转换成 YOLO 风格的 .txt。 - - 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 - -每张图输出一个同名 .txt，每行 `class_id cx cy w h`（相对宽高归一化）； -类别顺序默认按扫描到的名字排序，也可以用 `--classes scratch bubble ...` 固定。 -""" - -from __future__ import annotations - -import argparse -import json -import sys -import xml.etree.ElementTree as ET -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -from datasets.yolo import IMAGE_SUFFIXES, format_label_lines, write_data_yaml # noqa: E402 - - -def from_labelme(path: Path) -> dict: - data = json.loads(path.read_text(encoding="utf-8")) - boxes, labels = [], [] - for shape in data.get("shapes", []): - points = shape["points"] - xs = [p[0] for p in points] - ys = [p[1] for p in points] - boxes.append([min(xs), min(ys), max(xs), max(ys)]) - labels.append(shape["label"]) - return { - "image": Path(data.get("imagePath", path.stem + ".png")).name, - "width": data.get("imageWidth"), - "height": data.get("imageHeight"), - "boxes": boxes, - "labels": labels, - } - - -def from_voc(path: Path) -> dict: - root = ET.parse(path).getroot() - size = root.find("size") - boxes, labels = [], [] - for obj in root.findall("object"): - bnd = obj.find("bndbox") - if bnd is None: - continue - boxes.append([float(bnd.findtext(key, "0")) for key in ("xmin", "ymin", "xmax", "ymax")]) - labels.append(obj.findtext("name", "unknown")) - return { - "image": Path(root.findtext("filename") or (path.stem + ".jpg")).name, - "width": int(size.findtext("width", "0")) if size is not None else None, - "height": int(size.findtext("height", "0")) if size is not None else None, - "boxes": boxes, - "labels": labels, - } - - -PARSERS = {"labelme": (from_labelme, "*.json"), "voc": (from_voc, "*.xml")} - - -def image_size(record: dict, images_dir: Path | None) -> tuple[int, int]: - """标注里没写尺寸就去读图片本身——归一化坐标必须要有宽高。""" - if record.get("width") and record.get("height"): - return int(record["width"]), int(record["height"]) - if images_dir is None: - raise ValueError(f"{record['image']} 的标注里没有图像尺寸，请用 --images 指定图片目录") - - from PIL import Image - - stem = Path(record["image"]).stem - for suffix in IMAGE_SUFFIXES: - candidate = images_dir / f"{stem}{suffix}" - if candidate.exists(): - with Image.open(candidate) as image: - return image.size - raise FileNotFoundError(f"在 {images_dir} 下找不到 {record['image']}") - - -def main() -> int: - parser = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - parser.add_argument("--format", choices=sorted(PARSERS), required=True) - parser.add_argument("--input", required=True, help="标注文件所在目录") - parser.add_argument("--output", required=True, help="输出的 labels 目录") - parser.add_argument("--images", default=None, help="图片目录；标注里没有尺寸时用来读宽高") - parser.add_argument("--classes", nargs="*", default=None, help="固定类别顺序（缺省按名字排序）") - parser.add_argument("--data-yaml", default=None, help="顺带写一份 YOLO data.yaml") - args = parser.parse_args() - - convert, pattern = PARSERS[args.format] - files = sorted(Path(args.input).rglob(pattern)) - if not files: - print(f"在 {args.input} 下没有找到 {pattern}") - return 1 - - records = [convert(file) for file in files] - classes = args.classes or sorted({label for r in records for label in r["labels"]}) - class_to_idx = {name: i for i, name in enumerate(classes)} - - unknown = {label for r in records for label in r["labels"] if label not in class_to_idx} - if unknown: - print(f"标注中出现了 --classes 未列出的类别: {sorted(unknown)}") - return 1 - - images_dir = Path(args.images) if args.images else None - out_dir = Path(args.output) - out_dir.mkdir(parents=True, exist_ok=True) - - total_boxes = 0 - for record in records: - width, height = image_size(record, images_dir) - lines = format_label_lines( - record["boxes"], [class_to_idx[label] for label in record["labels"]], width, height - ) - total_boxes += len(lines) - # 无缺陷图也写空文件，方便核对"标了但没有目标" vs "漏标" - (out_dir / f"{Path(record['image']).stem}.txt").write_text( - "\n".join(lines) + ("\n" if lines else ""), encoding="utf-8" - ) - - print(f"已转换 {len(records)} 个文件、{total_boxes} 个框 -> {out_dir}") - print(f"类别顺序（class_id 即下标）: {classes}") - if args.data_yaml: - path = write_data_yaml(args.data_yaml, classes, root=out_dir.parent) - print(f"data.yaml 已写入 {path}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/export_onnx.py b/scripts/export_onnx.py deleted file mode 100755 index d2dd042..0000000 --- a/scripts/export_onnx.py +++ /dev/null @@ -1,77 +0,0 @@ -#!/usr/bin/env python -"""导出 ONNX，供实时设备（TensorRT / OpenVINO）部署。 - - python scripts/export_onnx.py --config configs/train/baseline.yaml \ - --checkpoint runs/baseline/best.pth --output runs/baseline/model.onnx - -导出在 CPU 上完成，不需要 GPU；需要额外依赖： - uv sync --group export -""" - -from __future__ import annotations - -import argparse -import sys -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -import torch # noqa: E402 - -from datasets import load_classes # noqa: E402 -from engine import load_checkpoint # noqa: E402 -from models import build_model # noqa: E402 -from utils import load_config # noqa: E402 - - -def main() -> int: - parser = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - parser.add_argument("--config", required=True) - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--output", required=True) - parser.add_argument("--opset", type=int, default=17) - parser.add_argument("--dynamic-batch", action="store_true") - args = parser.parse_args() - - cfg = load_config(args.config) - model_cfg = dict(cfg["model"]) - model_cfg.setdefault("num_classes", len(load_classes(cfg["classes"]))) - model_cfg.setdefault("label_mode", str(cfg.get("label_mode", "multilabel"))) - model = build_model(model_cfg) - load_checkpoint(args.checkpoint, model, map_location="cpu") - model.eval() - - height, width = cfg.get("image_size", [512, 512]) - channels = 1 if cfg.get("gray") else 3 - dummy = torch.randn(1, channels, int(height), int(width)) - - output_names = ["heatmap", "wh", "offset"] + (["logits"] if model.cls_head is not None else []) - dynamic_axes = {"images": {0: "batch"}} if args.dynamic_batch else None - if dynamic_axes: - dynamic_axes.update({name: {0: "batch"} for name in output_names}) - - try: - import onnxscript # noqa: F401 - except ImportError: - print("缺少导出依赖，请先执行: uv sync --group export", file=sys.stderr) - return 1 - - out_path = Path(args.output) - out_path.parent.mkdir(parents=True, exist_ok=True) - torch.onnx.export( - model, - dummy, - str(out_path), - input_names=["images"], - output_names=output_names, - opset_version=args.opset, - dynamic_axes=dynamic_axes, - ) - print(f"已导出 {out_path}（输入 {tuple(dummy.shape)}，输出 stride={model.out_stride}）") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/generate_boxes.py b/scripts/generate_boxes.py index dfc593f..0da49f7 100644 --- a/scripts/generate_boxes.py +++ b/scripts/generate_boxes.py @@ -154,8 +154,9 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--stage", type=int, default=6, - help="gradcam 的目标层（4~7 = layer1~4）。默认 6 即 layer3（8×8）；" - "layer4 只有 4×4 太粗，layer2 响应太散", + help="gradcam 的目标层（torchvision 骨架：4~7 = layer1~4）。默认 6 即 layer3（8×8）；" + "layer4 只有 4×4 太粗，layer2 响应太散。-1 表示用骨架最终输出——" + "tiny_conv 这类不是 nn.Sequential 的骨架只能这样取", ) parser.add_argument("--max-evals", type=int, default=500, help="shap 每张图的前向次数") parser.add_argument("--smooth-radius", type=int, default=2, help="shap 归因图平滑半径，0=不平滑") @@ -211,6 +212,19 @@ def main() -> int: for param in model.parameters(): # 归因需要梯度，但不更新权重 param.requires_grad_(True) + stage = args.stage if args.stage >= 0 else None + if stage is not None and not isinstance( + getattr(model.backbone, "features", None), torch.nn.Sequential + ): + # tiny_conv 这类骨架没有 nn.Sequential 的 features，取不了中间层； + # 退回最终输出而不是直接崩掉 + logger.warning( + "骨架 %s 不是 nn.Sequential，无法按 stage 取中间层，改用最终输出" + "（分辨率更低，如需更精细请换 torchvision 骨架）", + type(model.backbone).__name__, + ) + stage = None + min_area_ratio = ( args.min_area_ratio if args.min_area_ratio is not None else DEFAULT_MIN_AREA_RATIO[args.method] @@ -253,7 +267,7 @@ def main() -> int: predicted = model(images).argmax(1) targets = predicted if args.target == "predicted" else labels agree += int((predicted == labels).sum()) - heats = grad_cam(model.backbone, model.head, images, targets, stage=args.stage) + heats = grad_cam(model.backbone, model.head, images, targets, stage=stage) heat_shape = tuple(heats.shape[-2:]) for heat, target, meta in zip( @@ -310,7 +324,7 @@ def main() -> int: "heatmap_shape": list(heat_shape) if heat_shape else None, "seconds_per_image": (time.time() - started) / max(total, 1), "box_params": box_kwargs, - "stage": args.stage if args.method == "gradcam" else None, + "stage": stage if args.method == "gradcam" else None, "max_evals": args.max_evals if args.method == "shap" else None, "prediction_matches_folder_label": agree / total if total else 0.0, "boxes_per_image": {str(k): v for k, v in sorted(box_counts.items())}, diff --git a/scripts/profile_inference.py b/scripts/profile_inference.py deleted file mode 100755 index 4d59e9a..0000000 --- a/scripts/profile_inference.py +++ /dev/null @@ -1,218 +0,0 @@ -#!/usr/bin/env python -"""推理性能与显存分析（torch.profiler）。 - -部署在实时设备上，显存占用和单帧耗时是硬指标，因此这个脚本在推理路径上跑 -`torch.profiler`，给出算子级的显存与耗时排行，外加 CUDA 显存峰值统计。 - - # 用配置里的输入尺寸，batch=1 跑 50 次 - python scripts/profile_inference.py --config configs/train/baseline.yaml \ - --checkpoint runs/baseline/best.pth - - # 扫不同 batch，并导出 chrome trace（chrome://tracing 或 https://ui.perfetto.dev 打开） - python scripts/profile_inference.py --config configs/train/baseline.yaml \ - --batch-size 1 4 8 --trace runs/baseline/trace.json - -无 GPU 时自动退回 CPU（只有耗时与 CPU 内存，没有显存数据）。 -""" - -from __future__ import annotations - -import argparse -import json -import sys -import time -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -import torch # noqa: E402 -from torch.profiler import ProfilerActivity, profile, record_function # noqa: E402 - -from datasets import load_classes # noqa: E402 -from models import build_model, decode_heatmap # noqa: E402 -from utils import load_config, setup_logging # noqa: E402 -from utils.device import ( # noqa: E402 - DeviceNotAvailableError, - PeakMemory, - resolve_device, - synchronize, -) - -logger = setup_logging() - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - parser.add_argument("--config", required=True) - parser.add_argument("--checkpoint", default=None, help="不给则用随机初始化的权重，只看结构开销") - parser.add_argument("--batch-size", type=int, nargs="+", default=[1]) - parser.add_argument("--image-size", type=int, nargs=2, default=None, help="默认取配置的 image_size") - parser.add_argument("--iters", type=int, default=50, help="计时轮数") - parser.add_argument("--warmup", type=int, default=10, help="预热轮数（不计入统计）") - parser.add_argument("--device", default=None, help="auto / cuda / mps / cpu，默认 auto") - parser.add_argument("--amp", action="store_true", help="用 float16 autocast 推理") - parser.add_argument("--postprocess", action="store_true", help="把 heatmap 解码也计入") - parser.add_argument("--rows", type=int, default=15, help="算子排行显示条数") - parser.add_argument("--trace", default=None, help="导出 chrome trace 的路径") - parser.add_argument("--out", default=None, help="把汇总结果写入 json") - return parser.parse_args() - - -def build(args: argparse.Namespace, device: torch.device) -> torch.nn.Module: - cfg = load_config(args.config) - model_cfg = dict(cfg["model"]) - model_cfg.setdefault("num_classes", len(load_classes(cfg["classes"]))) - model_cfg.setdefault("label_mode", str(cfg.get("label_mode", "multilabel"))) - model = build_model(model_cfg) - - if args.checkpoint: - from engine import load_checkpoint - - load_checkpoint(args.checkpoint, model, map_location="cpu") - model.to(device).eval() - model.cfg = cfg - return model - - -@torch.no_grad() -def run_once(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespace) -> None: - device_type = images.device.type - with torch.autocast(device_type, dtype=torch.float16, enabled=args.amp and device_type != "cpu"): - with record_function("forward"): - outputs = model(images) - if args.postprocess: - with record_function("postprocess"): - decode_heatmap( - outputs["heatmap"].float(), - outputs["wh"].float(), - outputs["offset"].float(), - stride=model.out_stride, - ) - - -def measure(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespace) -> dict: - """预热后计时，并统计 CUDA 显存峰值。""" - device = images.device - peak = PeakMemory(device) - - for _ in range(args.warmup): - run_once(model, images, args) - synchronize(device) - peak.reset() - - start = time.perf_counter() - for _ in range(args.iters): - run_once(model, images, args) - synchronize(device) - elapsed = time.perf_counter() - start - - batch = images.shape[0] - result = { - "batch_size": batch, - "latency_ms": elapsed / args.iters * 1000, - "throughput_fps": args.iters * batch / elapsed, - } - if device.type != "cpu": - result.update( - { - "peak_allocated_mb": peak.peak_mb, - "peak_reserved_mb": peak.reserved_mb, - "peak_is_exact": peak.exact, - "weights_mb": sum(p.numel() * p.element_size() for p in model.parameters()) / 1024**2, - } - ) - return result - - -def profile_once(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespace) -> None: - """算子级 profiling：按显存占用与耗时各排一次序。""" - is_cuda = images.device.type == "cuda" - # torch.profiler 没有 MPS 的 activity，Metal 侧的 kernel 抓不到，只能看 CPU 侧下发耗时 - activities = [ProfilerActivity.CPU] + ([ProfilerActivity.CUDA] if is_cuda else []) - - with profile(activities=activities, profile_memory=True, record_shapes=True) as prof: - for _ in range(max(args.iters // 5, 1)): - run_once(model, images, args) - synchronize(images.device) - - memory_key = "self_cuda_memory_usage" if is_cuda else "self_cpu_memory_usage" - time_key = "self_cuda_time_total" if is_cuda else "self_cpu_time_total" - print(f"\n===== 显存占用 Top {args.rows}（按 {memory_key}）=====") - print(prof.key_averages().table(sort_by=memory_key, row_limit=args.rows)) - print(f"\n===== 耗时 Top {args.rows}（按 {time_key}）=====") - print(prof.key_averages().table(sort_by=time_key, row_limit=args.rows)) - - if args.trace: - trace_path = Path(args.trace) - trace_path.parent.mkdir(parents=True, exist_ok=True) - prof.export_chrome_trace(str(trace_path)) - print(f"\nchrome trace 已导出: {trace_path}（用 chrome://tracing 或 ui.perfetto.dev 打开）") - - -def main() -> int: - args = parse_args() - try: - device = resolve_device(args.device or "auto") - except DeviceNotAvailableError as exc: - logger.error("%s", exc) - return 1 - if device.type == "cpu": - logger.warning("当前在 CPU 上运行：没有显存数据，只能看耗时与 CPU 内存。") - elif device.type == "mps": - logger.warning( - "MPS 上 torch.profiler 抓不到 Metal kernel（没有对应的 ProfilerActivity），" - "算子明细只反映 CPU 侧下发耗时；显存用 torch.mps 采样，峰值是下界。" - ) - - model = build(args, device) - cfg = model.cfg - height, width = args.image_size or cfg.get("image_size", [512, 512]) - channels = 1 if cfg.get("gray") else 3 - - logger.info( - "模型 %s | 参数量 %.2fM | 输出 stride %d | 输入 %dx%dx%d | 设备 %s | amp=%s", - cfg["model"]["name"], - sum(p.numel() for p in model.parameters()) / 1e6, - model.out_stride, - channels, - height, - width, - device, - args.amp, - ) - - results = [] - for batch_size in args.batch_size: - images = torch.randn(batch_size, channels, int(height), int(width), device=device) - result = measure(model, images, args) - results.append(result) - memory = ( - f" | 显存峰值 allocated {result['peak_allocated_mb']:.1f} MB" - f" / reserved {result['peak_reserved_mb']:.1f} MB" - if "peak_allocated_mb" in result - else "" - ) - logger.info( - "batch=%-3d 单次 %.2f ms | %.1f FPS%s", - result["batch_size"], - result["latency_ms"], - result["throughput_fps"], - memory, - ) - - # 算子级明细只对最大的 batch 跑一次，避免输出过长。 - images = torch.randn(max(args.batch_size), channels, int(height), int(width), device=device) - profile_once(model, images, args) - - if args.out: - out_path = Path(args.out) - out_path.parent.mkdir(parents=True, exist_ok=True) - out_path.write_text(json.dumps(results, indent=2, ensure_ascii=False), encoding="utf-8") - logger.info("汇总已写入 %s", out_path) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/split_dataset.py b/scripts/split_dataset.py deleted file mode 100755 index 81c38a1..0000000 --- a/scripts/split_dataset.py +++ /dev/null @@ -1,104 +0,0 @@ -#!/usr/bin/env python -"""按比例切分成 YOLO 的 train / val / test 目录结构。 - - 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 - -默认建软链（不复制数据），产出： - -``` -images/train/ images/val/ -labels/train/ labels/val/ -``` - -原始文件保留在 `images/` `labels/` 根目录下不动；重复执行会先清掉旧的划分目录。 -""" - -from __future__ import annotations - -import argparse -import random -import shutil -import sys -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -from datasets.yolo import IMAGE_SUFFIXES, load_classes, write_data_yaml # noqa: E402 - - -def main() -> int: - parser = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - parser.add_argument("--images", required=True, help="图片目录（划分前的平铺目录）") - parser.add_argument("--labels", default=None, help="标注目录，缺省为 images 同级的 labels") - parser.add_argument("--val-ratio", type=float, default=0.2) - parser.add_argument("--test-ratio", type=float, default=0.0) - parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--copy", action="store_true", help="复制文件而不是建软链") - parser.add_argument("--data-yaml", default=None, help="顺带更新/生成 data.yaml") - parser.add_argument("--classes", nargs="*", default=None, help="写 data.yaml 时的类别顺序") - args = parser.parse_args() - - images_dir = Path(args.images) - labels_dir = Path(args.labels) if args.labels else images_dir.parent / "labels" - splits = ("train", "val", "test") - - files = sorted( - p - for p in images_dir.iterdir() - if p.is_file() and p.suffix.lower() in IMAGE_SUFFIXES and p.parent.name not in splits - ) - if not files: - print(f"{images_dir} 下没有待划分的图片（子目录里的已划分文件会被跳过）") - return 1 - - random.Random(args.seed).shuffle(files) - n_val = int(len(files) * args.val_ratio) - n_test = int(len(files) * args.test_ratio) - partition = { - "val": files[:n_val], - "test": files[n_val : n_val + n_test], - "train": files[n_val + n_test :], - } - - for name, items in partition.items(): - img_out, label_out = images_dir / name, labels_dir / name - for directory in (img_out, label_out): - if directory.exists(): - shutil.rmtree(directory) - if not items: - continue - img_out.mkdir(parents=True, exist_ok=True) - label_out.mkdir(parents=True, exist_ok=True) - - empty = 0 - for image in items: - label = labels_dir / f"{image.stem}.txt" - place(image, img_out / image.name, args.copy) - if label.exists(): - place(label, label_out / label.name, args.copy) - empty += int(not label.read_text(encoding="utf-8").strip()) - else: - empty += 1 # 缺 .txt 等价于无缺陷 - print(f"{name}: {len(items)} 张（其中无缺陷 {empty} 张） -> {img_out}") - - if args.data_yaml: - classes = load_classes(args.classes or args.data_yaml) - path = write_data_yaml(args.data_yaml, classes, root=images_dir.parent) - print(f"data.yaml 已更新 {path}") - return 0 - - -def place(src: Path, dst: Path, copy: bool) -> None: - if dst.exists() or dst.is_symlink(): - dst.unlink() - if copy: - shutil.copy2(src, dst) - else: - dst.symlink_to(src.resolve()) - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tests/test_config.py b/tests/test_config.py index 34a51e9..2a2f6af 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -66,10 +66,13 @@ def test_load_config_detects_cycle(tmp_path): def test_project_configs_are_loadable(): - cfg = load_config("configs/train/baseline.yaml") - assert cfg["model"]["name"] == "defect_net" - assert cfg["data"]["train"]["dataset"]["name"] == "film_defect" - assert len(cfg["classes"]) > 0 + cluster = load_config("configs/cluster/film.yaml") + assert cluster["dataset"]["name"] == "unlabeled" + assert cluster["num_clusters"] > 0 + + classify = load_config("configs/classify/film.yaml") + assert classify["model"]["name"] == "classifier" + assert classify["train"]["monitor"] == "macro_f1" def test_replace_key_swaps_a_whole_subdict(tmp_path): @@ -116,46 +119,39 @@ def test_registry_reports_incompatible_fields_actionably(): def test_every_shipped_config_builds_its_components(config_path): """仓库里每个配置都必须能真的构建出组件。 - 回归测试：configs/train/resnet18.yaml 曾经无法使用——它同时继承 tiny_conv 与 - resnet18 的模型配置，合并后 backbone 混了两边的字段，构建时直接报错。 + 回归测试：派生配置曾经无法使用——它同时继承两份骨架配置， + 合并后 backbone 混了两边的字段（tiny_conv 的 widths 混进 torchvision），构建时直接报错。 """ - torch = pytest.importorskip("torch") - from datasets.yolo import load_classes - from models import build_model + pytest.importorskip("torch") from models.registry import BACKBONES cfg = load_config(config_path) - if "model" in cfg: - if "classes" not in cfg: - # configs/model/*.yaml 是片段，只能被 _base_ 组合，单独用没有 classes - pytest.skip(f"{config_path} 是配置片段，不含 classes") - model_cfg = dict(cfg["model"]) - model_cfg.setdefault("num_classes", len(load_classes(cfg["classes"]))) - # 单测里不下载 ImageNet 权重 + # 聚类配置：pretrained 用顶层 backbone，simclr 用 simclr.backbone + for backbone_cfg in (cfg.get("backbone"), cfg.get_path("simclr.backbone")): + if not backbone_cfg: + continue + backbone_cfg = dict(backbone_cfg) + if backbone_cfg.get("name") == "torchvision": + backbone_cfg["pretrained"] = False # 单测里不下载 ImageNet 权重 + backbone = BACKBONES.build(backbone_cfg) + assert backbone.out_channels > 0 + + # 分类配置 + model_cfg = cfg.get("model") + if model_cfg: + import torch + + from models import build_model + + model_cfg = dict(model_cfg) + model_cfg.setdefault("num_classes", 3) + model_cfg["init_from"] = None if model_cfg.get("backbone", {}).get("name") == "torchvision": model_cfg["backbone"] = {**model_cfg["backbone"], "pretrained": False} model = build_model(model_cfg) with torch.no_grad(): - out = model(torch.randn(1, 3, 64, 64)) - assert "heatmap" in out - - simclr_backbone = cfg.get_path("simclr.backbone") - if simclr_backbone: - backbone_cfg = dict(simclr_backbone) - if backbone_cfg.get("name") == "torchvision": - backbone_cfg["pretrained"] = False - backbone = BACKBONES.build(backbone_cfg) - assert backbone.out_channels > 0 - - -def test_eval_config_matches_its_training_backbone(): - """评估配置构建出的骨架必须与训练配置一致，否则权重根本加载不上。""" - for train_cfg, eval_cfg in [ - ("configs/train/baseline.yaml", "configs/eval/baseline.yaml"), - ("configs/train/resnet18.yaml", "configs/eval/resnet18.yaml"), - ]: - assert load_config(train_cfg)["model"] == load_config(eval_cfg)["model"], eval_cfg + assert model(torch.randn(1, 3, 64, 64)).shape[1] == model_cfg["num_classes"] @pytest.mark.parametrize( @@ -169,13 +165,13 @@ def test_eval_config_matches_its_training_backbone(): def test_set_rejects_typos_with_suggestion(override, expected_hint): """--set 打错字母会凭空造出没人读的键，训练照跑但改动无效——必须拦下并给出建议。""" with pytest.raises(KeyError) as excinfo: - load_config("configs/cluster/film_simclr_resnet50.yaml", [override]) + load_config("configs/cluster/film_simclr_resnet18.yaml", [override]) assert expected_hint in str(excinfo.value) def test_set_accepts_existing_keys(): cfg = load_config( - "configs/cluster/film_simclr_resnet50.yaml", + "configs/cluster/film_simclr_resnet18.yaml", ["simclr.optimizer.lr=0.001", "simclr.optimizer.weight_decay=0.0", "num_clusters=7"], ) assert cfg["simclr"]["optimizer"]["lr"] == 0.001 diff --git a/tests/test_dataset.py b/tests/test_dataset.py deleted file mode 100644 index f1d2ac3..0000000 --- a/tests/test_dataset.py +++ /dev/null @@ -1,151 +0,0 @@ -"""YOLO 风格标注的数据集读取。""" - -import pytest - -pytest.importorskip("torch") -pytest.importorskip("PIL") - -from pathlib import Path # noqa: E402 - -import numpy as np # noqa: E402 -import torch # noqa: E402 -from PIL import Image # noqa: E402 - -from datasets import FilmDefectDataset, collate_fn # noqa: E402 - -CLASSES = ["scratch", "bubble"] -TRANSFORMS = [ - {"name": "Resize", "size": [64, 64]}, - {"name": "ToTensor"}, - {"name": "Normalize", "mean": [0.5, 0.5, 0.5], "std": [0.5, 0.5, 0.5]}, -] - - -@pytest.fixture -def dataset_root(tmp_path): - """YOLO 目录：images/ 与同层 labels/，128x128 的图，标注为归一化 cx cy w h。""" - images = tmp_path / "images" - labels = tmp_path / "labels" - images.mkdir() - labels.mkdir() - - for name in ("a.png", "b.png"): - Image.fromarray(np.zeros((128, 128, 3), dtype=np.uint8)).save(images / name) - # 像素框 (10,10,50,50) -> 中心 (30,30)/128, 宽高 40/128 - (labels / "a.txt").write_text("0 0.234375 0.234375 0.312500 0.312500\n", encoding="utf-8") - (labels / "b.txt").write_text("", encoding="utf-8") # 空文件 = 无缺陷 - return images, labels - - -def build(dataset_root, **kwargs): - images, _ = dataset_root - return FilmDefectDataset( - img_dir=images, classes=CLASSES, transforms=TRANSFORMS, stride=4, **kwargs - ) - - -def test_label_dir_follows_yolo_convention(dataset_root): - images, labels = dataset_root - dataset = build(dataset_root) - assert len(dataset) == 2 - assert dataset.label_file(images / "a.png") == labels / "a.txt" - - -def test_sample_shapes_and_box_denormalization(dataset_root): - dataset = build(dataset_root) - sample = dataset[0] - - assert sample["image"].shape == (3, 64, 64) - assert sample["target"]["heatmap"].shape == (2, 16, 16) - # 归一化框 -> 128px 像素框 (10,10,50,50) -> Resize 到 64 后减半 - assert sample["boxes"][0].tolist() == pytest.approx([5.0, 5.0, 25.0, 25.0], abs=1e-3) - assert sample["labels"].tolist() == [0] - assert sample["image_label"].tolist() == [1.0, 0.0] - - -def test_empty_label_file_is_defect_free(dataset_root): - sample = build(dataset_root)[1] - assert sample["boxes"].numel() == 0 - assert sample["target"]["mask"].sum() == 0 - assert sample["image_label"].sum() == 0 - - -def test_missing_label_file_is_defect_free(dataset_root): - _, labels = dataset_root - (labels / "b.txt").unlink() - assert build(dataset_root)[1]["boxes"].numel() == 0 - - -def test_explicit_label_dir(dataset_root, tmp_path): - images, labels = dataset_root - other = tmp_path / "yolo_labels" - other.mkdir() - (other / "a.txt").write_text("1 0.5 0.5 0.25 0.25\n", encoding="utf-8") - (other / "b.txt").write_text("", encoding="utf-8") - - dataset = build(dataset_root, label_dir=other) - assert dataset[0]["labels"].tolist() == [1] - - -def test_list_file_controls_order(dataset_root, tmp_path): - images, _ = dataset_root - list_file = tmp_path / "val.txt" - list_file.write_text("images/b.png\nimages/a.png\n", encoding="utf-8") - - dataset = build(dataset_root, list_file=list_file) - assert [Path(item).name for item in [dataset[i]["meta"]["path"] for i in range(2)]] == [ - "b.png", - "a.png", - ] - - -def test_collate_keeps_boxes_as_list(dataset_root): - dataset = build(dataset_root) - batch = collate_fn([dataset[0], dataset[1]]) - - assert batch["images"].shape == (2, 3, 64, 64) - assert batch["target"]["heatmap"].shape == (2, 2, 16, 16) - assert isinstance(batch["boxes"], list) and len(batch["boxes"]) == 2 - assert batch["image_labels"].shape == (2, 2) - - -def test_class_id_out_of_range_raises(dataset_root): - _, labels = dataset_root - (labels / "a.txt").write_text("7 0.5 0.5 0.2 0.2\n", encoding="utf-8") - with pytest.raises(ValueError, match="越界"): - _ = build(dataset_root)[0] - - -def test_gray_mode_yields_single_channel(dataset_root): - images, _ = dataset_root - transforms = [ - {"name": "Resize", "size": [64, 64]}, - {"name": "ToTensor"}, - {"name": "Normalize", "mean": [0.5], "std": [0.5]}, - ] - dataset = FilmDefectDataset( - img_dir=images, classes=CLASSES, transforms=transforms, stride=4, gray=True - ) - assert dataset[0]["image"].shape == (1, 64, 64) - - -def test_classes_from_data_yaml(dataset_root, tmp_path): - pytest.importorskip("yaml") - from datasets import write_data_yaml - - images, _ = dataset_root - data_yaml = write_data_yaml(tmp_path / "data.yaml", CLASSES, root=tmp_path) - dataset = FilmDefectDataset( - img_dir=images, classes=data_yaml, transforms=TRANSFORMS, stride=4 - ) - assert dataset.classes == CLASSES - - -def test_metrics_on_perfect_prediction(): - from engine.metrics import DetectionMetric - - metric = DetectionMetric(iou_threshold=0.5) - boxes = torch.tensor([[0.0, 0.0, 10.0, 10.0]]) - metric.update(boxes, torch.tensor([0]), boxes, torch.tensor([0])) - result = metric.compute() - assert result["precision"] == 1.0 and result["recall"] == 1.0 diff --git a/tests/test_generate_boxes.py b/tests/test_generate_boxes.py index cef7329..2e4d057 100644 --- a/tests/test_generate_boxes.py +++ b/tests/test_generate_boxes.py @@ -94,21 +94,25 @@ def test_both_methods_produce_a_yolo_dataset(project, tmp_path, monkeypatch, met assert all(0.0 <= float(v) <= 1.0 for v in parts[1:]) -def test_output_is_readable_by_the_detection_dataset(project, tmp_path, monkeypatch): - """产出要能直接喂给检测训练——这是这一步存在的意义。""" - from datasets import FilmDefectDataset +def test_output_is_a_wellformed_yolo_dataset(project, tmp_path, monkeypatch): + """产出是给外部检测训练用的交换格式，必须自洽： + data.yaml 的类别与 labels/ 对得上，每行标签可解析且坐标在 [0,1] 内。""" + import yaml out = tmp_path / "out" run(base_argv(project, out, "gradcam"), monkeypatch) - dataset = FilmDefectDataset( - img_dir=out / "images", classes=str(out / "data.yaml"), - transforms=[{"name": "Resize", "size": [64, 64]}, {"name": "ToTensor"}], stride=4, - ) - assert len(dataset) == 5 - sample = dataset[0] - assert sample["image"].shape == (3, 64, 64) - assert sample["boxes"].shape[1] == 4 + meta = yaml.safe_load((out / "data.yaml").read_text(encoding="utf-8")) + names = [meta["names"][i] for i in sorted(meta["names"])] + assert meta["nc"] == len(names) == 2 + + for label in (out / "labels").rglob("*.txt"): + assert (out / "images" / label.parent.name / f"{label.stem}.png").is_file() + for line in label.read_text(encoding="utf-8").strip().splitlines(): + class_id, cx, cy, w, h = line.split() + assert 0 <= int(class_id) < meta["nc"] + assert all(0.0 <= float(v) <= 1.0 for v in (cx, cy, w, h)) + assert float(w) > 0 and float(h) > 0 def test_images_are_hard_links_not_copies(project, tmp_path, monkeypatch): @@ -224,3 +228,59 @@ def test_duplicate_source_names_are_rejected(tmp_path): gb._emit(source=Path("/data/dirA/x.png"), **common) with pytest.raises(gb.NameCollision, match="prefix-parent"): gb._emit(source=Path("/data/dirB/x.png"), **common) # 同名，撞同一个标签路径 + + +def test_gradcam_falls_back_when_backbone_has_no_sequential_features(tmp_path, monkeypatch): + """tiny_conv 不是 nn.Sequential，取不了中间层——应当退回最终输出并警告，而不是崩掉。""" + from engine.checkpoint import save_checkpoint + from models import build_model + + root = tmp_path / "dataset" + (root / "a").mkdir(parents=True) + rng = np.random.default_rng(0) + for i in range(2): + Image.fromarray(rng.integers(0, 255, (40, 56, 3), dtype=np.uint8)).save( + root / "a" / f"a_{i}.png" + ) + + backbone = {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]} + checkpoint = tmp_path / "best.pth" + save_checkpoint(checkpoint, build_model( + {"name": "classifier", "num_classes": 1, "backbone": dict(backbone)} + )) + + config = tmp_path / "cfg.yaml" + config.write_text( + f""" +device: cpu +data: + root: {root} + classes: null + exclude: [] + val_transforms: [{{name: Resize, size: [64, 64]}}, {{name: ToTensor}}] + loader: {{batch_size: 2, num_workers: 0}} +model: + name: classifier + backbone: {{name: tiny_conv, widths: [8, 16], depths: [1, 1]}} + init_from: null +""", + encoding="utf-8", + ) + + out = tmp_path / "out" + monkeypatch.setattr( + sys, "argv", + ["generate_boxes.py", "--out", str(out), "--method", "gradcam", + "--checkpoint", str(checkpoint), "--config", str(config), "--no-heatmaps"], + ) + assert gb.main() == 0 # 不崩 + + summary = json.loads((out / "summary.json").read_text("utf-8")) + assert summary["stage"] is None # 已退回最终输出 + assert len(list((out / "labels").rglob("*.txt"))) == 2 + + +def test_stage_minus_one_selects_the_final_output(project, tmp_path, monkeypatch): + out = tmp_path / "out" + assert run(base_argv(project, out, "gradcam") + ["--stage", "-1"], monkeypatch) == 0 + assert json.loads((out / "summary.json").read_text("utf-8"))["stage"] is None diff --git a/tests/test_label_mode.py b/tests/test_label_mode.py deleted file mode 100644 index b38d9db..0000000 --- a/tests/test_label_mode.py +++ /dev/null @@ -1,125 +0,0 @@ -"""多类(multilabel) / 单类(multiclass) 由 yaml 的 label_mode 决定，两条路径都要能跑通。""" - -import pytest - -torch = pytest.importorskip("torch") -pytest.importorskip("PIL") - -import numpy as np # noqa: E402 -from PIL import Image # noqa: E402 - -from datasets import FilmDefectDataset, collate_fn # noqa: E402 -from engine.losses import DetectionLoss # noqa: E402 -from engine.metrics import SingleLabelMetric # noqa: E402 -from models import build_model # noqa: E402 -from models.defect_net import classification_outputs # noqa: E402 - -CLASSES = ["scratch", "bubble"] -TRANSFORMS = [{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}] - - -def make_model(label_mode: str, num_classes: int = 2): - return build_model( - { - "name": "defect_net", - "num_classes": num_classes, - "label_mode": label_mode, - "backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}, - "neck": {"name": "upsample", "channels": 8, "num_upsamples": 1}, - "head": {"hidden_channels": 8}, - } - ) - - -@pytest.fixture -def dataset_factory(tmp_path): - """YOLO 目录：a.png 上 scratch 1 个框、bubble 2 个框；b.png 无缺陷。""" - images = tmp_path / "images" - labels = tmp_path / "labels" - images.mkdir() - labels.mkdir() - for name in ("a.png", "b.png"): - Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)).save(images / name) - - (labels / "a.txt").write_text( - "0 0.1875 0.1875 0.25 0.25\n" # scratch - "1 0.5000 0.1875 0.25 0.25\n" # bubble - "1 0.2500 0.5000 0.25 0.25\n", # bubble - encoding="utf-8", - ) - (labels / "b.txt").write_text("", encoding="utf-8") - - def build(label_mode: str): - return FilmDefectDataset( - img_dir=images, - classes=CLASSES, - transforms=TRANSFORMS, - stride=2, - label_mode=label_mode, - ) - - return build - - -def test_multilabel_dataset_yields_multi_hot(dataset_factory): - dataset = dataset_factory("multilabel") - assert dataset[0]["image_label"].tolist() == [1.0, 1.0] - assert dataset[1]["image_label"].tolist() == [0.0, 0.0] - - -def test_multiclass_dataset_yields_exclusive_index(dataset_factory): - dataset = dataset_factory("multiclass") - # bubble 出现 2 次、scratch 1 次 -> 取 bubble(下标 1) - assert dataset[0]["image_label"].item() == 1 - assert dataset[0]["image_label"].dtype == torch.long - # 无缺陷 -> 背景下标 = len(classes) - assert dataset[1]["image_label"].item() == len(CLASSES) - - -def test_collate_shapes_differ_by_mode(dataset_factory): - multilabel = collate_fn([dataset_factory("multilabel")[i] for i in range(2)]) - multiclass = collate_fn([dataset_factory("multiclass")[i] for i in range(2)]) - assert multilabel["image_labels"].shape == (2, len(CLASSES)) - assert multiclass["image_labels"].shape == (2,) - # 定位监督与 label_mode 无关 - assert torch.equal(multilabel["target"]["heatmap"], multiclass["target"]["heatmap"]) - - -def test_classification_head_width(): - assert classification_outputs(5, "multilabel") == 5 - assert classification_outputs(5, "multiclass") == 6 - assert make_model("multilabel").cls_head.num_outputs == 2 - assert make_model("multiclass").cls_head.num_outputs == 3 - - -def test_invalid_label_mode_raises(): - with pytest.raises(ValueError, match="label_mode"): - make_model("single") - with pytest.raises(ValueError, match="label_mode"): - DetectionLoss(label_mode="single") - - -@pytest.mark.parametrize("label_mode", ["multilabel", "multiclass"]) -def test_end_to_end_loss_backward(dataset_factory, label_mode): - model = make_model(label_mode) - criterion = DetectionLoss(label_mode=label_mode) - dataset = dataset_factory(label_mode) - batch = collate_fn([dataset[0], dataset[1]]) - - losses = criterion(model(batch["images"]), batch) - losses["total"].backward() - - assert set(losses) == {"heatmap", "wh", "offset", "cls", "total"} - assert torch.isfinite(losses["total"]) - assert any(p.grad is not None and p.grad.abs().sum() > 0 for p in model.cls_head.parameters()) - - -def test_single_label_metric_counts_only_present_classes(): - metric = SingleLabelMetric(num_outputs=3) - logits = torch.tensor([[9.0, 0.0, 0.0], [0.0, 9.0, 0.0], [9.0, 0.0, 0.0]]) - metric.update(logits, torch.tensor([0, 1, 1])) - result = metric.compute() - - assert result["cls_accuracy"] == pytest.approx(2 / 3) - # 类别 2 没出现，不参与 macro 平均 - assert result["cls_f1"] == pytest.approx((2 / 3 + 2 / 3) / 2) diff --git a/tests/test_models.py b/tests/test_models.py index 3123050..1685e3b 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1,96 +1,12 @@ +"""卷积骨架与分类模型。""" + import pytest torch = pytest.importorskip("torch") from models import build_model # noqa: E402 from models.backbones.resnet import _SUPPORTED as _SUPPORTED_ARCHS # noqa: E402 -from models.postprocess import decode_heatmap # noqa: E402 - - -@pytest.fixture -def model_cfg(): - return { - "name": "defect_net", - "num_classes": 3, - "backbone": {"name": "tiny_conv", "widths": [8, 16, 32], "depths": [1, 1, 1]}, - "neck": {"name": "upsample", "channels": 16, "num_upsamples": 2}, - "head": {"hidden_channels": 16}, - } - - -def test_forward_shapes_and_stride(model_cfg): - model = build_model(model_cfg).eval() - # stem(2) * 2 个下采样 stage = 8，再上采样 4 倍 => stride 2 - assert model.out_stride == 2 - - with torch.no_grad(): - out = model(torch.randn(2, 3, 64, 64)) - - size = 64 // model.out_stride - assert out["heatmap"].shape == (2, 3, size, size) - assert out["wh"].shape == (2, 2, size, size) - assert out["offset"].shape == (2, 2, size, size) - assert out["logits"].shape == (2, 3) - - -def test_backbone_is_convolution_only(model_cfg): - model = build_model(model_cfg) - forbidden = (torch.nn.MultiheadAttention, torch.nn.Transformer, torch.nn.LSTM, torch.nn.GRU) - assert not any(isinstance(m, forbidden) for m in model.modules()) - - -def test_without_classification_head(model_cfg): - model = build_model({**model_cfg, "with_classification": False}).eval() - with torch.no_grad(): - out = model(torch.randn(1, 3, 64, 64)) - assert "logits" not in out - - -def test_decode_heatmap_recovers_planted_box(): - heatmap = torch.full((1, 2, 32, 32), -10.0) - heatmap[0, 1, 8, 10] = 10.0 - wh = torch.zeros(1, 2, 32, 32) - wh[0, :, 8, 10] = torch.tensor([4.0, 6.0]) - offset = torch.zeros(1, 2, 32, 32) - - (result,) = decode_heatmap(heatmap, wh, offset, stride=4, topk=5, score_threshold=0.5) - - assert result["labels"].tolist() == [1] - assert result["boxes"][0].tolist() == pytest.approx([32.0, 20.0, 48.0, 44.0]) - - -def test_losses_decrease_on_a_single_batch(model_cfg): - from engine.losses import DetectionLoss - - torch.manual_seed(0) - model = build_model(model_cfg) - criterion = DetectionLoss() - optimizer = torch.optim.Adam(model.parameters(), lr=1e-2) - - images = torch.randn(2, 3, 64, 64) - size = 64 // model.out_stride - target = { - "heatmap": torch.zeros(2, 3, size, size), - "wh": torch.zeros(2, 2, size, size), - "offset": torch.zeros(2, 2, size, size), - "mask": torch.zeros(2, size, size), - } - target["heatmap"][0, 1, 5, 5] = 1.0 - target["mask"][0, 5, 5] = 1.0 - target["wh"][0, :, 5, 5] = 3.0 - batch = {"target": target, "image_labels": torch.zeros(2, 3)} - - first = last = None - for step in range(5): - losses = criterion(model(images), batch) - optimizer.zero_grad() - losses["total"].backward() - optimizer.step() - total = float(losses["total"].detach()) - first = total if step == 0 else first - last = total - - assert last < first +from models.heads.projection import ProjectionHead # noqa: E402 @pytest.mark.parametrize("arch", sorted(_SUPPORTED_ARCHS)) @@ -118,10 +34,23 @@ def test_unsupported_arch_is_rejected(): BACKBONES.build({"name": "torchvision", "arch": "vit_b_16"}) +def test_tiny_conv_shapes(): + from models.registry import BACKBONES + + backbone = BACKBONES.build( + {"name": "tiny_conv", "widths": [8, 16, 32], "depths": [1, 1, 1]} + ).eval() + with torch.no_grad(): + feat = backbone(torch.randn(2, 3, 64, 64)) + assert feat.shape[:2] == (2, 32) + assert backbone.out_channels == 32 + # stem 下采样 2 倍，其后每个 stage(除第一个)各下采样 2 倍 + assert backbone.out_stride == 8 + + @pytest.mark.parametrize("arch", ["resnet18", "shufflenet_v2_x1_0", "mobilenet_v3_small"]) def test_alternative_backbones_drive_simclr(arch): """换骨架后 SimCLR 的投影头要能接上（out_channels 必须准）。""" - from models.heads.projection import ProjectionHead from models.registry import BACKBONES backbone = BACKBONES.build({"name": "torchvision", "arch": arch, "pretrained": False}).eval() @@ -130,18 +59,13 @@ def test_alternative_backbones_drive_simclr(arch): assert head(backbone(torch.randn(2, 3, 64, 64))).shape == (2, 16) -def test_alternative_backbone_drives_defect_net(): - """检测模型换骨架同样要能跑通（neck 的 in_channels 由 backbone 决定）。""" - model = build_model( - { - "name": "defect_net", - "num_classes": 3, - "backbone": {"name": "torchvision", "arch": "shufflenet_v2_x1_0", "pretrained": False}, - "neck": {"name": "upsample", "channels": 16, "num_upsamples": 3}, - "head": {"hidden_channels": 16}, - } - ).eval() - assert model.out_stride == 4 +def test_classifier_shapes_and_convolution_only(): + model = build_model({ + "name": "classifier", "num_classes": 4, + "backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}, + "classification": {"dropout": 0.1}, + }).eval() with torch.no_grad(): - out = model(torch.randn(1, 3, 64, 64)) - assert out["heatmap"].shape == (1, 3, 16, 16) + assert model(torch.randn(3, 3, 64, 64)).shape == (3, 4) + forbidden = (torch.nn.MultiheadAttention, torch.nn.Transformer, torch.nn.LSTM) + assert not any(isinstance(m, forbidden) for m in model.modules()) diff --git a/tests/test_profile_inference.py b/tests/test_profile_inference.py deleted file mode 100644 index 474656b..0000000 --- a/tests/test_profile_inference.py +++ /dev/null @@ -1,70 +0,0 @@ -"""推理 profiler 脚本：CPU 上也要能跑出耗时结果（有 GPU 时额外给出显存峰值）。""" - -import argparse -import importlib.util -from pathlib import Path - -import pytest - -torch = pytest.importorskip("torch") - -from models import build_model # noqa: E402 - -SCRIPT = Path(__file__).resolve().parents[1] / "scripts" / "profile_inference.py" - - -@pytest.fixture(scope="module") -def script(): - spec = importlib.util.spec_from_file_location("profile_inference", SCRIPT) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -@pytest.fixture -def model(): - return build_model( - { - "name": "defect_net", - "num_classes": 2, - "backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}, - "neck": {"name": "upsample", "channels": 8, "num_upsamples": 1}, - "head": {"hidden_channels": 8}, - } - ).eval() - - -@pytest.fixture -def args(): - return argparse.Namespace(iters=2, warmup=1, amp=False, postprocess=True, rows=3, trace=None) - - -def test_measure_reports_latency_and_memory(script, model, args): - images = torch.randn(2, 3, 32, 32) - result = script.measure(model, images, args) - - assert result["batch_size"] == 2 - assert result["latency_ms"] > 0 - assert result["throughput_fps"] > 0 - if torch.cuda.is_available(): - assert result["peak_allocated_mb"] > 0 - assert result["peak_reserved_mb"] >= result["peak_allocated_mb"] - else: - assert "peak_allocated_mb" not in result - - -def test_profile_exports_trace(script, model, args, tmp_path, capsys): - args.trace = str(tmp_path / "trace.json") - script.profile_once(model, torch.randn(1, 3, 32, 32), args) - - assert Path(args.trace).exists() - out = capsys.readouterr().out - assert "显存占用 Top" in out and "耗时 Top" in out - - -def test_postprocess_toggle_runs_both_paths(script, model, args): - images = torch.randn(1, 3, 32, 32) - args.postprocess = False - script.run_once(model, images, args) - args.postprocess = True - script.run_once(model, images, args) diff --git a/tests/test_target_encoder.py b/tests/test_target_encoder.py deleted file mode 100644 index 9209674..0000000 --- a/tests/test_target_encoder.py +++ /dev/null @@ -1,42 +0,0 @@ -import pytest - -pytest.importorskip("numpy") -pytest.importorskip("torch") - -import numpy as np # noqa: E402 - -from datasets.target_encoder import CenterTargetEncoder, gaussian_radius # noqa: E402 - - -def test_gaussian_radius_is_positive_and_scales(): - assert gaussian_radius(40, 40) > gaussian_radius(10, 10) > 0 - - -def test_encoder_marks_single_center(): - encoder = CenterTargetEncoder(num_classes=3, stride=4) - boxes = np.array([[40.0, 40.0, 80.0, 80.0]]) - labels = np.array([1]) - - target = encoder(boxes, labels, (256, 256)) - - assert target["heatmap"].shape == (3, 64, 64) - assert target["mask"].sum() == 1 - # 中心 (60, 60) / stride 4 = (15, 15) - assert target["mask"][15, 15] == 1 - assert target["heatmap"][1].max() == pytest.approx(1.0) - assert target["heatmap"][0].max() == 0 - assert target["wh"][:, 15, 15].tolist() == pytest.approx([10.0, 10.0]) - - -def test_encoder_skips_invalid_and_out_of_range_boxes(): - encoder = CenterTargetEncoder(num_classes=2, stride=4) - boxes = np.array([[10.0, 10.0, 5.0, 5.0], [900.0, 900.0, 950.0, 950.0]]) - target = encoder(boxes, np.array([0, 1]), (256, 256)) - assert target["mask"].sum() == 0 - - -def test_encoder_respects_max_objects(): - encoder = CenterTargetEncoder(num_classes=1, stride=4, max_objects=2) - boxes = np.array([[i * 20.0, 0.0, i * 20.0 + 16, 16.0] for i in range(5)]) - target = encoder(boxes, np.zeros(5, dtype=np.int64), (256, 256)) - assert target["mask"].sum() == 2 diff --git a/tests/test_train_smoke.py b/tests/test_train_smoke.py deleted file mode 100644 index 12dc778..0000000 --- a/tests/test_train_smoke.py +++ /dev/null @@ -1,217 +0,0 @@ -"""train.py / eval.py 的端到端冒烟：YOLO 数据目录 -> 训练 -> 评估。 - -覆盖入口脚本的接线：classes 解析、label_mode 下发、设备解析。固定用 CPU 跑， -免得测试受机器上有没有 GPU 影响。 -""" - -import sys - -import pytest - -pytest.importorskip("torch") -pytest.importorskip("yaml") -pytest.importorskip("PIL") - -import numpy as np # noqa: E402 -import torch # noqa: E402 -import yaml # noqa: E402 -from PIL import Image # noqa: E402 - -import eval as eval_mod # noqa: E402 -import train as train_mod # noqa: E402 -from datasets import write_data_yaml # noqa: E402 - -CLASSES = ["scratch", "bubble"] -TRANSFORMS = [ - {"name": "Resize", "size": [32, 32]}, - {"name": "ToTensor"}, - {"name": "Normalize", "mean": [0.5, 0.5, 0.5], "std": [0.5, 0.5, 0.5]}, -] - - -@pytest.fixture -def project(tmp_path): - """造一个最小的 YOLO 数据集 + 配置，classes 用 data.yaml 的路径给。""" - for split in ("train", "val"): - images = tmp_path / "images" / split - labels = tmp_path / "labels" / split - images.mkdir(parents=True) - labels.mkdir(parents=True) - for i in range(2): - Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)).save(images / f"{i}.png") - content = "0 0.5 0.5 0.25 0.25\n" if i else "" - (labels / f"{i}.txt").write_text(content, encoding="utf-8") - - data_yaml = write_data_yaml(tmp_path / "data.yaml", CLASSES, root=tmp_path) - - def dataset(split): - return { - "dataset": {"name": "film_defect", "img_dir": str(tmp_path / "images" / split), - "transforms": TRANSFORMS}, - "loader": {"batch_size": 2, "num_workers": 0, "shuffle": split == "train"}, - } - - config = { - "classes": str(data_yaml), - "label_mode": "multilabel", - "image_size": [32, 32], - "device": "cpu", - "seed": 0, - "model": { - "name": "defect_net", - "backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}, - "neck": {"name": "upsample", "channels": 8, "num_upsamples": 1}, - "head": {"hidden_channels": 8}, - }, - "data": {"train": dataset("train"), "val": dataset("val")}, - "train": { - "epochs": 1, - "amp": False, - "log_interval": 10, - "optimizer": {"name": "adamw", "lr": 0.001}, - "swanlab": {"enabled": False}, - }, - } - config_path = tmp_path / "config.yaml" - config_path.write_text(yaml.safe_dump(config, allow_unicode=True), encoding="utf-8") - return config_path, tmp_path - - -def run(module, argv, monkeypatch): - monkeypatch.setattr(sys, "argv", argv) - return module.main() - - -def test_train_then_eval(project, monkeypatch, tmp_path): - config_path, root = project - work_dir = tmp_path / "runs" - - assert run(train_mod, ["train.py", "--config", str(config_path), "--work-dir", str(work_dir)], - monkeypatch) == 0 - assert (work_dir / "best.pth").exists() - assert (work_dir / "config.yaml").exists() - - metrics_path = tmp_path / "metrics.json" - assert run( - eval_mod, - ["eval.py", "--config", str(config_path), "--checkpoint", str(work_dir / "best.pth"), - "--set", "device=cpu", "--out", str(metrics_path)], - monkeypatch, - ) == 0 - assert metrics_path.exists() - - -def test_classes_from_data_yaml_sizes_the_heads(project, monkeypatch, tmp_path): - """classes 写成 data.yaml 路径时，num_classes 必须是 2，而不是路径字符串的长度。""" - config_path, _ = project - work_dir = tmp_path / "runs2" - captured = {} - - real_build = train_mod.build_model - - def spy(cfg): - captured["cfg"] = dict(cfg) - return real_build(cfg) - - monkeypatch.setattr(train_mod, "build_model", spy) - - run(train_mod, ["train.py", "--config", str(config_path), "--work-dir", str(work_dir)], monkeypatch) - - assert captured["cfg"]["num_classes"] == len(CLASSES) - assert captured["cfg"]["label_mode"] == "multilabel" - - -def test_unavailable_device_fails_loudly(project, monkeypatch, tmp_path): - """显式请求机器上没有的后端时必须退出，而不是悄悄退回 CPU。""" - from utils.device import available_backends - - missing = next((b for b in ("cuda", "mps") if b not in available_backends()), None) - if missing is None: - pytest.skip("这台机器 cuda 与 mps 都可用") - - assert run( - train_mod, - ["train.py", "--config", str(config_path_of(project)), "--work-dir", str(tmp_path / "runs3"), - "--device", missing], - monkeypatch, - ) == 1 - - -def test_device_flag_overrides_config(project, monkeypatch, tmp_path): - """--device 优先于配置里的 device。""" - config_path, _ = project - work_dir = tmp_path / "runs4" - assert run( - train_mod, - ["train.py", "--config", str(config_path), "--work-dir", str(work_dir), "--device", "cpu"], - monkeypatch, - ) == 0 - assert (work_dir / "best.pth").exists() - - -def config_path_of(project): - return project[0] - - -def test_resume_takes_hyperparams_from_config(project, monkeypatch, tmp_path): - """检测训练续训：lr / weight_decay / 总轮数以 yaml 为准，检查点只提供进度。 - - 回归测试：load_checkpoint 直接 optimizer.load_state_dict，把超参连同动量一起还原， - 改了配置再续训会被静默覆盖；而且此前根本没保存调度器状态，学习率调度会从头开始。 - """ - import math - - import yaml as yaml_mod - - config_path, root = project - cfg = yaml_mod.safe_load(config_path.read_text(encoding="utf-8")) - cfg["train"].update({"epochs": 4, "optimizer": {"name": "adamw", "lr": 1e-3, "weight_decay": 1e-4}, - "scheduler": {"name": "cosine"}}) - config_path.write_text(yaml_mod.safe_dump(cfg, allow_unicode=True), encoding="utf-8") - - work_dir = tmp_path / "runs" - assert run(train_mod, ["train.py", "--config", str(config_path), "--work-dir", str(work_dir)], - monkeypatch) == 0 - - captured = {} - real_fit = train_mod.Trainer.fit - - def probe(self): - captured["groups"] = [dict(g) for g in self.optimizer.param_groups] - captured["scheduler"] = self.scheduler - captured["start_epoch"] = self.start_epoch - return 0.0 - - monkeypatch.setattr(train_mod.Trainer, "fit", probe) - assert run( - train_mod, - ["train.py", "--config", str(config_path), "--work-dir", str(work_dir), - "--resume", str(work_dir / "last.pth"), "--set", - "train.epochs=10", "train.optimizer.lr=5e-4", "train.optimizer.weight_decay=0.05"], - monkeypatch, - ) == 0 - monkeypatch.setattr(train_mod.Trainer, "fit", real_fit) - - groups = captured["groups"] - assert groups[0]["initial_lr"] == pytest.approx(5e-4) # 不是检查点里的 1e-3 - assert groups[0]["weight_decay"] == pytest.approx(0.05) # 不是检查点里的 1e-4 - assert groups[1]["weight_decay"] == 0.0 # BN/bias 分组结构没被冲掉 - - scheduler = captured["scheduler"] - assert scheduler.T_max == 10 - assert scheduler.last_epoch == captured["start_epoch"] == 4 # 调度进度已快进，不是从 0 开始 - assert groups[0]["lr"] == pytest.approx(5e-4 * (1 + math.cos(math.pi * 4 / 10)) / 2) - - -def test_checkpoint_carries_scheduler_and_progress(project, monkeypatch, tmp_path): - """检查点要带上调度器状态与 global_step，否则续训无从恢复调度进度。""" - config_path, _ = project - work_dir = tmp_path / "runs" - assert run(train_mod, ["train.py", "--config", str(config_path), "--work-dir", str(work_dir)], - monkeypatch) == 0 - - state = torch.load(work_dir / "last.pth", map_location="cpu", weights_only=False) - expected = {"model", "optimizer", "scheduler", "epoch", "global_step", "best_metric", "train_cfg"} - assert expected <= set(state) - assert state["global_step"] > 0 - assert state["train_cfg"]["optimizer"]["lr"] is not None diff --git a/tests/test_yolo_format.py b/tests/test_yolo_format.py deleted file mode 100644 index fa9a797..0000000 --- a/tests/test_yolo_format.py +++ /dev/null @@ -1,118 +0,0 @@ -"""YOLO 标注格式的解析与写出。""" - -from pathlib import Path - -import pytest - -pytest.importorskip("numpy") - -import numpy as np # noqa: E402 - -from datasets.yolo import ( # noqa: E402 - denormalize, - format_label_lines, - label_path_for, - list_images, - load_classes, - parse_label_file, - write_data_yaml, -) - - -def test_label_path_swaps_images_for_labels(): - img_dir = Path("/data/film/images/train") - assert label_path_for(img_dir / "0001.png", img_dir) == Path("/data/film/labels/train/0001.txt") - - -def test_label_path_keeps_subdirs_and_honours_explicit_dir(): - img_dir = Path("/data/film/images/train") - assert label_path_for(img_dir / "cam1/0001.png", img_dir) == Path( - "/data/film/labels/train/cam1/0001.txt" - ) - assert label_path_for(img_dir / "0001.png", img_dir, Path("/tmp/lab")) == Path("/tmp/lab/0001.txt") - - -def test_parse_center_format(tmp_path): - path = tmp_path / "a.txt" - path.write_text("0 0.5 0.5 0.4 0.2\n# 注释\n\n1 0.25 0.25 0.1 0.1\n", encoding="utf-8") - - boxes, labels = parse_label_file(path) - - assert labels.tolist() == [0, 1] - assert boxes[0].tolist() == pytest.approx([0.3, 0.4, 0.7, 0.6]) - assert denormalize(boxes, 100, 200)[0].tolist() == pytest.approx([30.0, 80.0, 70.0, 120.0]) - - -def test_missing_file_means_no_defect(tmp_path): - boxes, labels = parse_label_file(tmp_path / "nope.txt") - assert boxes.shape == (0, 4) and labels.shape == (0,) - - -def test_polygon_rows_fall_back_to_bounding_box(tmp_path): - path = tmp_path / "a.txt" - path.write_text("2 0.1 0.1 0.4 0.2 0.3 0.5\n", encoding="utf-8") - - boxes, labels = parse_label_file(path) - - assert labels.tolist() == [2] - assert boxes[0].tolist() == pytest.approx([0.1, 0.1, 0.4, 0.5]) - - -def test_coordinates_are_clamped_and_degenerate_rows_dropped(tmp_path): - path = tmp_path / "a.txt" - path.write_text("0 0.5 0.5 2.0 2.0\n0 0.5 0.5 0.0 0.0\n", encoding="utf-8") - - boxes, labels = parse_label_file(path) - - assert len(labels) == 1 - assert boxes[0].tolist() == pytest.approx([0.0, 0.0, 1.0, 1.0]) - - -def test_malformed_rows_report_file_and_line(tmp_path): - path = tmp_path / "a.txt" - path.write_text("0 0.5 0.5 0.4\n", encoding="utf-8") - with pytest.raises(ValueError, match=r"a\.txt:1"): - parse_label_file(path) - - path.write_text("scratch 0.5 0.5 0.4 0.2\n", encoding="utf-8") - with pytest.raises(ValueError, match="不是合法的 YOLO 标注行"): - parse_label_file(path) - - -def test_class_id_range_is_checked(tmp_path): - path = tmp_path / "a.txt" - path.write_text("3 0.5 0.5 0.4 0.2\n", encoding="utf-8") - with pytest.raises(ValueError, match="越界"): - parse_label_file(path, num_classes=2) - - -def test_format_lines_roundtrips_pixel_boxes(tmp_path): - boxes = np.array([[10.0, 20.0, 50.0, 60.0]]) - lines = format_label_lines(boxes, [1], width=100, height=200) - assert lines == ["1 0.300000 0.200000 0.400000 0.200000"] - - path = tmp_path / "a.txt" - path.write_text("\n".join(lines), encoding="utf-8") - parsed, labels = parse_label_file(path) - assert labels.tolist() == [1] - assert denormalize(parsed, 100, 200)[0].tolist() == pytest.approx([10.0, 20.0, 50.0, 60.0]) - - -def test_data_yaml_roundtrip(tmp_path): - pytest.importorskip("yaml") - path = write_data_yaml(tmp_path / "data.yaml", ["scratch", "bubble"], root=tmp_path) - assert load_classes(path) == ["scratch", "bubble"] - assert load_classes(["a", "b"]) == ["a", "b"] - - -def test_list_images_supports_dir_and_list_file(tmp_path): - images = tmp_path / "images" - images.mkdir() - for name in ("b.png", "a.jpg", "notes.txt"): - (images / name).write_bytes(b"") - - assert [p.name for p in list_images(images)] == ["a.jpg", "b.png"] - - list_file = tmp_path / "train.txt" - list_file.write_text("images/b.png\nimages/a.jpg\n", encoding="utf-8") - assert [p.name for p in list_images(images, list_file)] == ["b.png", "a.jpg"] diff --git a/train.py b/train.py deleted file mode 100755 index 2d5ea6c..0000000 --- a/train.py +++ /dev/null @@ -1,99 +0,0 @@ -#!/usr/bin/env python -"""训练入口。 - -用法: - python train.py --config configs/train/baseline.yaml - python train.py --config configs/train/baseline.yaml --set train.epochs=10 train.optimizer.lr=5e-4 - python train.py --config configs/train/baseline.yaml --device mps - -设备由配置的 `device` 决定，默认 `auto`：按 cuda > mps > cpu 挑最快的可用后端。 -显式指定但不可用时直接报错，不会悄悄退回 CPU。 -""" - -from __future__ import annotations - -import argparse -import sys -from pathlib import Path - -from datasets import build_dataloader, build_dataset, load_classes -from engine import Trainer -from models import build_model -from utils import check_label_mode, load_config, seed_everything, setup_logging -from utils.device import DeviceNotAvailableError, describe_device, resolve_device - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="薄膜缺陷分类定位模型训练") - parser.add_argument("--config", required=True, help="yaml 配置路径") - parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE", help="覆盖配置项") - parser.add_argument("--work-dir", default=None, help="覆盖配置中的 work_dir") - parser.add_argument("--device", default=None, help="覆盖配置中的 device：auto / cuda / mps / cpu") - parser.add_argument("--resume", default=None, help="从检查点恢复训练") - return parser.parse_args() - - -def main() -> int: - args = parse_args() - cfg = load_config(args.config, args.set) - - work_dir = Path(args.work_dir or cfg.get("work_dir", "runs/default")) - work_dir.mkdir(parents=True, exist_ok=True) - logger = setup_logging(work_dir / "train.log") - - try: - device = resolve_device(args.device or cfg.get("device", "auto")) - except DeviceNotAvailableError as exc: - logger.error("%s", exc) - return 1 - logger.info("训练设备: %s", describe_device(device)) - - seed_everything(int(cfg.get("seed", 42)), bool(cfg.get("deterministic", False))) - - classes = load_classes(cfg["classes"]) - # label_mode 由 yaml 决定，模型与数据集必须一致，这里统一下发。 - label_mode = check_label_mode(str(cfg.get("label_mode", "multilabel"))) - model_cfg = dict(cfg["model"]) - model_cfg.setdefault("num_classes", len(classes)) - model_cfg.setdefault("label_mode", label_mode) - model = build_model(model_cfg).to(device) - logger.info( - "模型 %s 输出 stride=%d，参数量 %.2fM", - model_cfg["name"], - model.out_stride, - sum(p.numel() for p in model.parameters()) / 1e6, - ) - - loaders = {} - for split in ("train", "val"): - split_cfg = cfg["data"].get(split) - if not split_cfg: - loaders[split] = None - continue - dataset = build_dataset( - split_cfg["dataset"], classes=classes, stride=model.out_stride, label_mode=label_mode - ) - loaders[split] = build_dataloader(dataset, **dict(split_cfg.get("loader", {}))) - logger.info("%s 集样本数: %d", split, len(dataset)) - - trainer = Trainer( - model=model, - train_loader=loaders["train"], - val_loader=loaders["val"], - cfg=cfg.get("train", {}), - device=device, - work_dir=work_dir, - full_config=cfg.to_dict(), - ) - - if args.resume: - trainer.resume(args.resume) - - cfg.dump(work_dir / "config.yaml") - best = trainer.fit() - logger.info("训练结束，最佳 %s=%.4f", trainer.monitor, best) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/utils/__init__.py b/utils/__init__.py index 008b1c8..d7cad96 100644 --- a/utils/__init__.py +++ b/utils/__init__.py @@ -10,7 +10,6 @@ from utils.device import ( resolve_amp, resolve_device, ) -from utils.labels import LABEL_MODES, MULTICLASS, MULTILABEL, check_label_mode from utils.logging import get_logger, setup_logging from utils.registry import Registry from utils.seed import seed_everything @@ -25,10 +24,6 @@ __all__ = [ "memory_stats", "PeakMemory", "DeviceNotAvailableError", - "LABEL_MODES", - "MULTILABEL", - "MULTICLASS", - "check_label_mode", "get_logger", "setup_logging", "Registry", diff --git a/utils/labels.py b/utils/labels.py deleted file mode 100644 index 753bd33..0000000 --- a/utils/labels.py +++ /dev/null @@ -1,13 +0,0 @@ -"""图像级标签形式的共用定义（模型与数据集都要用到，放这里避免互相 import）。""" - -from __future__ import annotations - -MULTILABEL = "multilabel" -MULTICLASS = "multiclass" -LABEL_MODES = (MULTILABEL, MULTICLASS) - - -def check_label_mode(label_mode: str) -> str: - if label_mode not in LABEL_MODES: - raise ValueError(f"未知 label_mode: {label_mode}，可选 {LABEL_MODES}") - return label_mode