File size: 1,830 Bytes
6e802a4 | 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 | # test_level_decay.py: The per-codebook depth loss weighting.
import pathlib
import sys
import pytest
import torch
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent))
FISH_CB = [4096] + [1024] * 9
def _mk(decay):
from model import DepthModule
torch.manual_seed(0)
return DepthModule(d_model=32, num_levels=10, codebook_size=1024,
d_depth=64, num_layers=2, cross_frame=0,
codebook_sizes=FISH_CB, level_decay=decay)
def _weights(m, semantic_weight=1.0):
d = m.level_decay
ws = [semantic_weight if k == 0 else d ** (k - 1) for k in range(m.num_levels)]
if d != 1.0:
ws = [w * m.num_levels / sum(ws) for w in ws]
return ws
def test_default_is_flat_and_unchanged():
m = _mk(1.0)
assert m.level_decay == 1.0
assert _weights(m)[1:] == [1.0] * 9
def test_decay_emphasises_coarse_codebooks():
ws = _weights(_mk(0.8))[1:]
assert all(a > b for a, b in zip(ws, ws[1:])), ws
def test_decay_preserves_total_scale():
for d in (0.6, 0.8, 0.95):
assert sum(_weights(_mk(d))) == pytest.approx(10.0, rel=1e-6)
@pytest.mark.parametrize("decay", [1.0, 0.8])
def test_forward_loss_finite_and_differentiable(decay):
m = _mk(decay)
h = torch.randn(4, 32)
tgt = torch.stack([torch.randint(0, s, (4,)) for s in FISH_CB], dim=1)
loss = m.forward_loss(h, tgt, semantic_weight=1.5)
assert torch.isfinite(loss)
loss.backward()
assert m.h_proj.weight.grad is not None
def test_decay_actually_changes_the_loss():
torch.manual_seed(1)
h = torch.randn(4, 32)
tgt = torch.stack([torch.randint(0, s, (4,)) for s in FISH_CB], dim=1)
flat = _mk(1.0).forward_loss(h, tgt).item()
dec = _mk(0.7).forward_loss(h, tgt).item()
assert abs(flat - dec) > 1e-4, (flat, dec)
|