Spaces:
Sleeping
Sleeping
File size: 3,662 Bytes
dc8d52b 2d52135 dc8d52b | 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 | import pytest
import torch
from mini_transformer.modules.feedforward import FeedForwardLayer
# =======================
# Constructor checks
# =======================
def test_ctor_type_checks():
with pytest.raises(TypeError):
FeedForwardLayer("64", 256, 0.1) # d_model
with pytest.raises(TypeError):
FeedForwardLayer(64, "256", 0.1) # d_ff
with pytest.raises(TypeError):
FeedForwardLayer(64, 256, "0.1") # dropout_rate
def test_ctor_value_checks():
with pytest.raises(ValueError):
FeedForwardLayer(0, 256, 0.1)
with pytest.raises(ValueError):
FeedForwardLayer(64, 0, 0.1)
with pytest.raises(ValueError):
FeedForwardLayer(64, 256, -1.0)
with pytest.raises(ValueError):
FeedForwardLayer(64, 256, 1.0) # upper bound excluded
def test_ctor_happy_path_defaults_and_attrs():
ffn = FeedForwardLayer(64, 256) # default dropout=0.1
assert ffn.d_model == 64 and ffn.d_ff == 256
assert isinstance(ffn.dropout, torch.nn.Dropout)
assert abs(ffn.dropout.p - 0.1) < 1e-9
# =======================
# Forward: input validation
# =======================
def test_forward_type_and_rank_checks():
ffn = FeedForwardLayer(32, 64, 0.1)
with pytest.raises(TypeError):
ffn("not a tensor")
with pytest.raises(ValueError):
ffn(torch.randn(5, 32)) # rank 2
with pytest.raises(ValueError):
ffn(torch.randn(2, 3, 16)) # D mismatch (expects 32)
# =======================
# Forward: shapes, device, dtype, zero-length
# =======================
@pytest.mark.parametrize("B,S,D", [(1, 1, 16), (2, 7, 32), (3, 0, 24)])
def test_forward_shapes_device_dtype_and_zero_len(B, S, D):
device = (
torch.device(f"cuda:{torch.cuda.current_device()}")
if torch.cuda.is_available()
else torch.device("cpu")
)
ffn = FeedForwardLayer(D, D * 2, 0.2).to(device)
x = torch.randn(B, S, D, device=device, dtype=torch.float32)
y = ffn(x)
assert y.shape == (B, S, D)
assert y.device == device
assert y.dtype == x.dtype
# =======================
# Forward: gradient flow
# =======================
def test_gradients_flow():
D, H = 32, 64
ffn = FeedForwardLayer(D, H, 0.1)
x = torch.randn(2, 5, D, requires_grad=True)
y = ffn(x)
loss = y.pow(2).mean()
loss.backward()
assert x.grad is not None and torch.isfinite(x.grad).all()
# =======================
# Dropout behavior
# =======================
def test_dropout_zero_equals_no_dropout_path():
D, H = 16, 32
# Build two modules: one with p=0.0, one with p>0 but switched to eval()
ffn0 = FeedForwardLayer(D, H, 0.0)
ffnx = FeedForwardLayer(D, H, 0.5).eval() # eval disables dropout
ffnx.load_state_dict(ffn0.state_dict(), strict=False)
x = torch.randn(2, 4, D)
y0 = ffn0(x)
yx = ffnx(x)
# With dropout disabled both should be equal (same weights distrib not identical, but same computation tree)
assert torch.allclose(y0, yx, atol=1e-6, rtol=1e-6)
def test_dropout_training_changes_output_vs_eval():
D, H = 32, 64
ffn = FeedForwardLayer(D, H, 0.5)
x = torch.randn(3, 6, D)
torch.manual_seed(123)
ffn.train()
y_train = ffn(x)
torch.manual_seed(123)
ffn.eval()
y_eval = ffn(x)
# Same seed but eval disables dropout -> outputs should differ
assert not torch.allclose(y_train, y_eval)
def test_dropout_is_noop_on_zero_length():
D, H = 32, 64
ffn = FeedForwardLayer(D, H, 0.7).train()
x = torch.randn(2, 0, D) # zero-length sequence
y = ffn(x)
assert y.shape == (2, 0, D) # no crash; shape preserved
|