| import math
|
| import pytest
|
| import torch
|
| from glint.box import (BoxHead, BoxLossWeights, box_answer, box_iou, box_loss, box_metrics, decode_box, decode_edges,
|
| dfl_loss, generalized_iou, max_prob_confidence, order_box, two_hot)
|
|
|
| K = 64
|
|
|
|
|
| def logits_from_target(target, bins=K):
|
| return two_hot(target, bins).clamp_min(1e-12).log()
|
|
|
|
|
|
|
| @pytest.mark.parametrize("bins", [2, 16, 64])
|
| def test_two_hot_round_trip(bins):
|
| torch.manual_seed(0)
|
| t = torch.rand(100, 4)
|
| t[0] = torch.tensor([0.0, 1.0, 0.0, 1.0])
|
| t[1] = torch.tensor([0.5, 1 / (bins - 1), 1 - 1e-7, min(1.0, 2 / (bins - 1))])
|
| th = two_hot(t, bins)
|
| assert th.shape == (100, 4, bins)
|
| assert torch.allclose(th.sum(-1), torch.ones(100, 4))
|
| assert ((th > 0).sum(-1) <= 2).all()
|
| assert torch.allclose(decode_edges(logits_from_target(t, bins)), t, atol=1e-5)
|
|
|
|
|
| def test_two_hot_clamps_out_of_range():
|
| t = torch.tensor([[-0.2, 1.3, 0.25, 0.75]])
|
| assert torch.allclose(decode_edges(logits_from_target(t)), t.clamp(0, 1), atol=1e-5)
|
|
|
|
|
|
|
| def test_dfl_matches_manual_and_is_minimal_at_target():
|
| torch.manual_seed(0)
|
| t = torch.tensor([[0.1, 0.4, 0.33, 0.9]])
|
| z = torch.randn(1, 4, K, requires_grad=True)
|
| loss = dfl_loss(z, t)
|
| manual = 0.0
|
| for e in range(4):
|
| x = float(t[0, e]) * (K - 1); lo = int(math.floor(x)); w = x - lo
|
| lp = z[0, e].log_softmax(-1)
|
| manual += -((1 - w) * lp[lo] + w * lp[lo + 1])
|
| assert torch.allclose(loss, manual / 4, atol=1e-5)
|
| loss.backward()
|
| assert torch.isfinite(z.grad).all() and z.grad.abs().sum() > 0
|
|
|
| th = two_hot(t, K)
|
| ent = -(th * th.clamp_min(1e-12).log()).sum(-1).mean()
|
| assert torch.allclose(dfl_loss(logits_from_target(t), t), ent, atol=1e-4)
|
| assert dfl_loss(logits_from_target(t), t) < dfl_loss(torch.zeros(1, 4, K), t)
|
|
|
|
|
| def test_dfl_exact_bin_target_is_one_hot_ce():
|
| t = torch.full((1, 4), 10 / (K - 1))
|
| z = torch.randn(1, 4, K)
|
| assert torch.allclose(dfl_loss(z, t), -z.log_softmax(-1)[..., 10].mean(), atol=1e-5)
|
|
|
|
|
|
|
| def test_iou_giou_known_values():
|
| a = torch.tensor([0.0, 0.5, 0.0, 0.5])
|
| assert torch.allclose(box_iou(a, a), torch.tensor(1.0))
|
| assert torch.allclose(generalized_iou(a, a), torch.tensor(1.0))
|
| b = torch.tensor([0.25, 0.75, 0.25, 0.75])
|
| assert torch.allclose(box_iou(a, b), torch.tensor(0.0625 / (0.25 + 0.25 - 0.0625)), atol=1e-6)
|
| c = torch.tensor([0.5, 1.0, 0.5, 1.0])
|
| assert torch.allclose(box_iou(a, c), torch.tensor(0.0))
|
| assert torch.allclose(generalized_iou(a, c), torch.tensor(-0.5), atol=1e-6)
|
| d = torch.tensor([0.0, 0.1, 0.9, 1.0])
|
| assert generalized_iou(a, d) < 0 and generalized_iou(a, d) > -1
|
|
|
|
|
| def test_giou_gradient_for_disjoint_boxes():
|
| p = torch.tensor([0.0, 0.1, 0.0, 0.1], requires_grad=True)
|
| g = torch.tensor([0.8, 0.9, 0.8, 0.9])
|
| (1 - generalized_iou(p, g)).backward()
|
| assert p.grad.abs().sum() > 0
|
|
|
|
|
| def test_iou_broadcast():
|
| a = torch.rand(5, 4).sort(-1).values[:, [0, 1, 2, 3]]
|
| a = order_box(a)
|
| assert box_iou(a, a).shape == (5,)
|
| assert torch.allclose(box_iou(a, a), torch.ones(5), atol=1e-5)
|
|
|
|
|
|
|
| def test_order_box_swaps_inverted_edges():
|
| e = torch.tensor([[0.8, 0.2, 0.9, 0.1], [0.1, 0.2, 0.3, 0.4]])
|
| assert torch.equal(order_box(e), torch.tensor([[0.2, 0.8, 0.1, 0.9], [0.1, 0.2, 0.3, 0.4]]))
|
| z = logits_from_target(torch.tensor([[0.7, 0.3, 0.6, 0.2]]))
|
| t, b, l, r = decode_box(z)[0].tolist()
|
| assert t <= b and l <= r
|
|
|
|
|
| def test_box_answer_jev_shape_and_strict_order():
|
| a = box_answer(torch.tensor([0.7, 0.3, 0.5, 0.5]), 1.3)
|
| assert a["type"] == "box" and set(a["box"]) == {"top", "bottom", "left", "right"}
|
| assert a["box"]["top"] < a["box"]["bottom"] and a["box"]["left"] < a["box"]["right"]
|
| assert a["box"]["top"] == 0.3 and a["box"]["bottom"] == 0.7
|
| assert a["confidence"] == 1.0
|
| z = box_answer([1.0, 1.0, 0.0, 0.0], 0.5)["box"]
|
| assert 0 <= z["top"] < z["bottom"] <= 1 and 0 <= z["left"] < z["right"] <= 1
|
|
|
|
|
|
|
| def test_box_loss_parts_and_backward():
|
| torch.manual_seed(0)
|
| z = torch.randn(3, 4, K, requires_grad=True); c = torch.zeros(3, requires_grad=True)
|
| tgt = torch.tensor([[0.1, 0.2, 0.3, 0.6], [0.5, 0.9, 0.1, 0.2], [0.0, 1.0, 0.0, 1.0]])
|
| loss, parts = box_loss(z, c, tgt, BoxLossWeights())
|
| assert set(parts) == {"dfl", "l1", "giou", "conf", "iou"}
|
| loss.backward()
|
| assert torch.isfinite(z.grad).all() and z.grad.abs().sum() > 0 and c.grad.abs().sum() > 0
|
| w0, _ = box_loss(z, c, tgt, BoxLossWeights(dfl=0, l1=0, giou=0, conf=0))
|
| assert float(w0.detach()) == 0.0
|
| perfect, pp = box_loss(logits_from_target(tgt), torch.full((3,), 10.0), tgt)
|
| assert pp["l1"] < 1e-4 and pp["giou"] < 1e-3 and pp["iou"] > 0.999
|
|
|
|
|
| def test_max_prob_confidence_range():
|
| z = torch.randn(2, 4, K)
|
| c = max_prob_confidence(z)
|
| assert c.shape == (2,) and ((c > 0) & (c <= 1)).all()
|
| assert torch.allclose(max_prob_confidence(logits_from_target(torch.rand(2, 4))), torch.ones(2), atol=1e-5)
|
|
|
|
|
| def test_box_metrics():
|
| m = box_metrics([0.9, 0.6, 0.3, 0.8], [0.9, 0.7, 0.2, 0.8], [0.5] * 4)
|
| assert m["n_box_questions"] == 4
|
| assert m["mean_iou"] == pytest.approx(0.65)
|
| assert m["iou_ge_0.5"] == 0.75 and m["iou_ge_0.75"] == 0.5
|
| assert 0 <= m["ece_15bin_conf_vs_iou50"] <= 1 and "ece_15bin_maxprob_vs_iou50" in m
|
| assert box_metrics([], []) == {"n_box_questions": 0}
|
|
|
|
|
|
|
| class _Rec:
|
| def __init__(self, spans, coords=None):
|
| from glint.schema import EncodedBoxQuestion
|
| self.box_questions = tuple(EncodedBoxQuestion(f"q{i}", s) for i, s in enumerate(spans))
|
| self.image_coords = coords
|
|
|
|
|
| def test_box_head_forward_backward_random_hidden():
|
| torch.manual_seed(0)
|
| H, L = 48, 40
|
| head = BoxHead(hidden_size=H, width=32, layers=2, heads=4, feedforward=64, bins=K)
|
| hidden = torch.randn(2, L, H, requires_grad=True)
|
| mask = torch.ones(2, L, dtype=torch.long); mask[1, 30:] = 0
|
| coords = (tuple(range(2, 14)), tuple(((r + .5) / 3, (c + .5) / 4) for r in range(3) for c in range(4)))
|
| recs = [_Rec([(20, 25), (26, 30)], coords), _Rec([(5, 9)])]
|
| out = head(hidden, mask, recs)
|
| assert out[0].logits.shape == (2, 4, K) and out[0].conf_logit.shape == (2,)
|
| assert out[1].logits.shape == (1, 4, K) and out[0].question_ids == ("q0", "q1")
|
| tgt = torch.tensor([[0.1, 0.3, 0.2, 0.8], [0.6, 0.7, 0.1, 0.4]])
|
| loss, _ = box_loss(out[0].logits, out[0].conf_logit, tgt)
|
| loss2, _ = box_loss(out[1].logits, out[1].conf_logit, tgt[:1])
|
| (loss + loss2).backward()
|
| assert hidden.grad is not None and torch.isfinite(hidden.grad).all()
|
| assert hidden.grad[1, 30:].abs().sum() == 0
|
| for name, p in head.named_parameters():
|
| assert p.grad is not None, name
|
| assert head(hidden, mask, [_Rec([])])[0] is None
|
|
|
|
|
| def test_box_head_questions_do_not_interact():
|
| torch.manual_seed(0)
|
| head = BoxHead(hidden_size=16, width=16, layers=1, heads=2, feedforward=32, bins=8).eval()
|
| hidden = torch.randn(1, 20, 16); mask = torch.ones(1, 20, dtype=torch.long)
|
| both = head(hidden, mask, [_Rec([(2, 5), (8, 12)])])[0]
|
| one = head(hidden, mask, [_Rec([(2, 5)])])[0]
|
| assert torch.allclose(both.logits[0], one.logits[0], atol=1e-5)
|
|
|
|
|
| def test_box_head_can_overfit_and_bf16():
|
| torch.manual_seed(0)
|
| head = BoxHead(hidden_size=16, width=32, layers=1, heads=4, feedforward=64, bins=K)
|
| hidden = torch.randn(1, 12, 16); mask = torch.ones(1, 12, dtype=torch.long)
|
| rec = [_Rec([(3, 6)])]; tgt = torch.tensor([[0.42, 0.47, 0.18, 0.83]])
|
| opt = torch.optim.AdamW(head.parameters(), lr=3e-3)
|
| for _ in range(150):
|
| o = head(hidden, mask, rec)[0]; loss, parts = box_loss(o.logits, o.conf_logit, tgt)
|
| opt.zero_grad(); loss.backward(); opt.step()
|
| assert parts["iou"] > 0.9
|
| with torch.no_grad():
|
| assert float(torch.sigmoid(head(hidden, mask, rec)[0].conf_logit)) > 0.7
|
| hb = head.to(torch.bfloat16)
|
| ob = hb(hidden.bfloat16(), mask, rec)[0]
|
| assert ob.logits.dtype == torch.bfloat16 and box_iou(decode_box(ob.logits)[0], tgt[0]) > 0.8
|
|
|
|
|
| def test_box_head_edges_do_not_collapse():
|
|
|
| torch.manual_seed(1)
|
| head = BoxHead(hidden_size=16, width=32, layers=1, heads=4, feedforward=64, bins=K)
|
| hidden = torch.randn(4, 12, 16); mask = torch.ones(4, 12, dtype=torch.long)
|
| tgt = torch.tensor([[0.48, 0.54, 0.50, 0.86], [0.53, 0.59, 0.08, 0.92], [0.20, 0.30, 0.10, 0.40], [0.70, 0.95, 0.60, 0.90]])
|
| opt = torch.optim.AdamW(head.parameters(), lr=1e-3)
|
| for step in range(400):
|
| i = step % 4
|
| o = head(hidden[i:i + 1], mask[i:i + 1], [_Rec([(3, 6)])])[0]
|
| loss, _ = box_loss(o.logits, o.conf_logit, tgt[i:i + 1])
|
| opt.zero_grad(); loss.backward(); opt.step()
|
| with torch.no_grad():
|
| out = torch.stack([decode_box(head(hidden[i:i + 1], mask[i:i + 1], [_Rec([(3, 6)])])[0].logits)[0] for i in range(4)])
|
| assert box_iou(out, tgt).mean() > 0.5
|
|
|
