Shaun ZhouandClaude Opus 5
d6adea5b05
88 个测试覆盖配置继承、注册表、GPU 守卫、YOLO 格式解析、数据集、模型形状、 标签模式、聚类原语与 SimCLR 损失，以及 train.py / eval.py / 聚类脚本的入口接线。 含两条回归测试： - PCA 白化放大噪声方向导致簇塌陷（默认 whiten=False 的依据） - classes 写成 data.yaml 路径时 num_classes 必须是类别数而非路径字符串长度 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
126 lines
4.5 KiB
Python
126 lines
4.5 KiB
Python
"""多类(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)
|