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>
138 lines
4.5 KiB
Python
138 lines
4.5 KiB
Python
"""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
|