| import torch |
|
|
| from objectmodel_v1.losses import ObjectModelCriterion |
| from objectmodel_v1.model import build_model |
| from objectmodel_v1.postprocess import decode_predictions |
|
|
|
|
| def tiny_config(dense_aux=True): |
| return { |
| "model": { |
| "num_classes": 5, |
| "input_size": 128, |
| "stem_channels": 16, |
| "backbone_channels": [24, 32, 48, 64], |
| "backbone_depths": [1, 1, 1, 1], |
| "hidden_dim": 48, |
| "fpn_depth": 1, |
| "latent_count": 8, |
| "latent_pool_sizes": [4, 2, 1], |
| "latent_layers": 1, |
| "decoder_layers": 2, |
| "num_queries": 12, |
| "num_heads": 4, |
| "local_points": 2, |
| "dropout": 0.0, |
| "dense_aux": dense_aux, |
| }, |
| "loss": { |
| "cost_class": 2.0, |
| "cost_bbox": 5.0, |
| "cost_giou": 2.0, |
| "weight_class": 2.0, |
| "weight_bbox": 5.0, |
| "weight_giou": 2.0, |
| "weight_dense": 1.0, |
| "focal_alpha": 0.25, |
| "focal_gamma": 2.0, |
| "aux_weight": 1.0, |
| "dense_topk": 3, |
| }, |
| } |
|
|
|
|
| def test_forward_shapes_and_ranges(): |
| model = build_model(tiny_config()).eval() |
| with torch.no_grad(): |
| output = model(torch.randn(2, 3, 128, 128)) |
| assert output["pred_logits"].shape == (2, 12, 5) |
| assert output["pred_boxes"].shape == (2, 12, 4) |
| assert len(output["aux_outputs"]) == 1 |
| assert "dense_outputs" not in output |
| assert torch.all((output["pred_boxes"] >= 0) & (output["pred_boxes"] <= 1)) |
|
|
|
|
| def test_loss_backward_with_empty_target(): |
| config = tiny_config() |
| model = build_model(config).train() |
| criterion = ObjectModelCriterion(config) |
| output = model(torch.randn(2, 3, 128, 128)) |
| targets = [ |
| { |
| "boxes": torch.tensor([[0.5, 0.5, 0.25, 0.3], [0.2, 0.2, 0.1, 0.1]]), |
| "labels": torch.tensor([1, 3]), |
| }, |
| {"boxes": torch.empty(0, 4), "labels": torch.empty(0, dtype=torch.long)}, |
| ] |
| losses = criterion(output, targets) |
| assert all(torch.isfinite(value) for value in losses.values()) |
| losses["loss_total"].backward() |
| gradients = [parameter.grad for parameter in model.parameters() if parameter.grad is not None] |
| assert gradients |
| assert all(torch.isfinite(gradient).all() for gradient in gradients) |
|
|
|
|
| def test_decode_is_nms_free_top_k_filter(): |
| model = build_model(tiny_config(dense_aux=False)).eval() |
| with torch.no_grad(): |
| output = model(torch.randn(1, 3, 128, 128)) |
| result = decode_predictions(output, [(240, 320)], confidence=0.0, top_k=4)[0] |
| assert result["boxes"].shape == (4, 4) |
| assert result["scores"].shape == (4,) |
| assert torch.all(result["boxes"][:, [0, 2]] <= 320) |
| assert torch.all(result["boxes"][:, [1, 3]] <= 240) |
|
|
