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>
93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
import pytest
|
|
|
|
torch = pytest.importorskip("torch")
|
|
|
|
from models import build_model # 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
|