File size: 10,260 Bytes
e379257 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 | 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()
# ---------------------------------------------------------------- bins encode / decode
@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]) # exact endpoints
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() # mass only on two adjacent bins
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)
# ---------------------------------------------------------------- DFL
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
# loss at the two-hot target logits == entropy of the two-hot distribution (its minimum)
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)
# ---------------------------------------------------------------- IoU / GIoU
def test_iou_giou_known_values():
a = torch.tensor([0.0, 0.5, 0.0, 0.5]) # top, bottom, left, right
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]) # overlap 0.25x0.25
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]) # touching corner, hull = 1
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 # IoU alone would have zero gradient here
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)
# ---------------------------------------------------------------- ordering
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"] # degenerate at the image border
assert 0 <= z["top"] < z["bottom"] <= 1 and 0 <= z["left"] < z["right"] <= 1
# ---------------------------------------------------------------- loss / metrics
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}
# ---------------------------------------------------------------- head forward / backward
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 # padding is ignored
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 # confidence tracks IoU
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():
# regression: with a tiny edge embedding all 4 edges decoded to the same value (the L1-optimal mean of the targets)
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
|