test: 单元测试与端到端冒烟
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>
This commit is contained in:
co-authored by
Claude Opus 5
parent
f433e5cf59
commit
d6adea5b05
@@ -0,0 +1,137 @@
|
||||
"""train.py / eval.py 的端到端冒烟:YOLO 数据目录 -> 训练 -> 评估。
|
||||
|
||||
覆盖入口脚本的接线(classes 解析、label_mode 下发、GPU 守卫),
|
||||
所以在 CPU 上跑,`require_cuda` 用 monkeypatch 放行。
|
||||
"""
|
||||
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("torch")
|
||||
pytest.importorskip("yaml")
|
||||
pytest.importorskip("PIL")
|
||||
|
||||
import numpy as np # 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
|
||||
from utils.device import NoNvidiaGpuError # 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"
|
||||
monkeypatch.setattr(train_mod, "require_cuda", lambda: ["FAKE-GPU"])
|
||||
|
||||
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)
|
||||
monkeypatch.setattr(train_mod, "require_cuda", lambda: ["FAKE-GPU"])
|
||||
|
||||
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_training_is_blocked_without_nvidia_gpu(project, monkeypatch, tmp_path):
|
||||
config_path, _ = project
|
||||
|
||||
def refuse():
|
||||
raise NoNvidiaGpuError("没有 Nvidia GPU")
|
||||
|
||||
monkeypatch.setattr(train_mod, "require_cuda", refuse)
|
||||
assert run(
|
||||
train_mod,
|
||||
["train.py", "--config", str(config_path), "--work-dir", str(tmp_path / "runs3")],
|
||||
monkeypatch,
|
||||
) == 1
|
||||
Reference in New Issue
Block a user