Spaces:
Running
Running
File size: 2,623 Bytes
a5ec84d | 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 | """Stage 2.3 — Decoder. Hermetic tests on random features + synthetic captions."""
import pytest
import torch
from capit.config import config
from capit.losses import caption_loss
from capit.models.decoder import Decoder
V = 12
B = 3
L = config.encoded_size**2
# captions sorted by length descending; PAD=0 START=1 END=2
CAPTIONS = torch.tensor(
[
[1, 5, 6, 7, 8, 2],
[1, 5, 6, 2, 0, 0],
[1, 7, 2, 0, 0, 0],
]
)
LENGTHS = torch.tensor([6, 4, 3])
@pytest.fixture(autouse=True)
def _seed():
torch.manual_seed(config.seed)
def _features() -> torch.Tensor:
return torch.randn(B, L, config.encoder_dim)
def test_forward_shapes():
logits, alphas, decode_lengths = Decoder(vocab_size=V)(_features(), CAPTIONS, LENGTHS)
assert decode_lengths == [5, 3, 2]
assert logits.shape == (B, 5, V)
assert alphas.shape == (B, 5, L)
def test_no_nan():
logits, alphas, _ = Decoder(vocab_size=V)(_features(), CAPTIONS, LENGTHS)
assert not torch.isnan(logits).any()
assert not torch.isnan(alphas).any()
def test_shrinking_leaves_finished_rows_zero():
logits, alphas, decode_lengths = Decoder(vocab_size=V)(_features(), CAPTIONS, LENGTHS)
for i, dl in enumerate(decode_lengths):
assert torch.all(logits[i, dl:] == 0)
assert torch.all(alphas[i, dl:] == 0)
def test_padded_positions_contribute_zero_loss():
logits, alphas, decode_lengths = Decoder(vocab_size=V)(_features(), CAPTIONS, LENGTHS)
base = caption_loss(logits, alphas, CAPTIONS)
corrupted = logits.clone()
for i, dl in enumerate(decode_lengths):
corrupted[i, dl:] = 999.0
assert torch.isclose(base, caption_loss(corrupted, alphas, CAPTIONS))
def test_forward_rejects_unsorted_batch():
captions = torch.tensor([[1, 7, 2, 0, 0, 0], [1, 5, 6, 7, 8, 2]]) # ascending lengths
lengths = torch.tensor([3, 6])
with pytest.raises(ValueError):
Decoder(vocab_size=V)(torch.randn(2, L, config.encoder_dim), captions, lengths)
def test_gradient_reaches_decoder_params():
dec = Decoder(vocab_size=V)
logits, alphas, _ = dec(_features(), CAPTIONS, LENGTHS)
caption_loss(logits, alphas, CAPTIONS).backward()
for p in (dec.embedding.weight, dec.fc.weight, dec.f_beta.weight):
assert p.grad is not None and p.grad.abs().sum() > 0
def test_greedy_returns_tokens_and_restores_mode():
dec = Decoder(vocab_size=V)
dec.train()
out = dec.greedy(_features(), start_id=1, end_id=2, max_len=10)
assert len(out) == B
assert all(isinstance(seq, list) and 2 not in seq for seq in out)
assert dec.training
|