refactor: 按实际的五步工作流裁掉无关代码
实际工作流是：simclr 训练 -> 聚类 -> 人工归类 -> 分类训练 -> GradCAM/SHAP 生成
热力图与 box。据此移除不参与其中的检测/定位链路，共 31 个文件、约 2600 行：
- 模型：defect_net、postprocess、necks/、heads/center_head
- 数据：film_defect（YOLO 读取端）、target_encoder
- 引擎：trainer、evaluator、losses
- 入口：train.py、eval.py
- 脚本：convert_annotations、split_dataset、export_onnx、profile_inference
- 配置：configs/{train,eval,model,dataset}
- utils/labels（多类/单类标签模式，仅检测分类头使用）
共享模块按实际用量瘦身：engine/metrics 只留 AverageMeter（检测指标随之移除）；
datasets/yolo 只留写出侧（IMAGE_SUFFIXES 与 write_data_yaml）——第 5 步产出的
YOLO 标签是给外部检测训练用的交换格式，读取端不再需要。
裁剪过程中修掉一个真 bug：generate_boxes 的 --stage 默认 6（layer3）对 tiny_conv
这类非 nn.Sequential 骨架直接抛异常。现在自动退回骨架最终输出并警告，
另加 --stage -1 显式选择最终输出。
五步工作流已端到端实跑验证（simclr 1 轮 -> 聚类 -> 导出 -> 分类 1 轮 ->
gradcam 与 shap 两种方法各生成一遍）。
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
57446c2f10
commit
86c86b105c
@@ -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/<name>/`）：`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 预训练模型
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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]
|
||||
@@ -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]
|
||||
@@ -1,8 +0,0 @@
|
||||
_base_:
|
||||
- ./baseline.yaml
|
||||
- ../model/resnet18.yaml
|
||||
|
||||
work_dir: runs/resnet18
|
||||
train:
|
||||
optimizer:
|
||||
lr: 0.0005
|
||||
+5
-13
@@ -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",
|
||||
]
|
||||
|
||||
+2
-27
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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],
|
||||
}
|
||||
@@ -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),
|
||||
}
|
||||
+6
-148
@@ -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]，中心点格式：
|
||||
|
||||
```
|
||||
<class_id> <cx> <cy> <w> <h>
|
||||
```
|
||||
|
||||
也接受 YOLO 分割格式的多边形（`<class_id> 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,
|
||||
|
||||
+3
-7
@@ -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",
|
||||
|
||||
@@ -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
|
||||
@@ -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())
|
||||
+4
-155
@@ -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()),
|
||||
}
|
||||
|
||||
@@ -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
|
||||
@@ -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())
|
||||
+4
-9
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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
|
||||
@@ -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"]
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
@@ -1,3 +0,0 @@
|
||||
from models.necks.upsample import UpsampleNeck
|
||||
|
||||
__all__ = ["UpsampleNeck"]
|
||||
@@ -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
|
||||
@@ -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
|
||||
+1
-2
@@ -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"]
|
||||
|
||||
@@ -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())
|
||||
@@ -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())
|
||||
@@ -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())},
|
||||
|
||||
@@ -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())
|
||||
@@ -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())
|
||||
+33
-37
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
+26
-102
@@ -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())
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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"]
|
||||
@@ -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())
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user