Spaces:
Sleeping
Sleeping
File size: 2,761 Bytes
dc8d52b 2d52135 dc8d52b 2d52135 dc8d52b 2d52135 dc8d52b 6edd10d dc8d52b 6edd10d 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 | import pytest
import torch
from mini_transformer.modules.encoder import EncoderLayer, TransformerEncoder
def _dev():
return (
torch.device(f"cuda:{torch.cuda.current_device()}")
if torch.cuda.is_available()
else torch.device("cpu")
)
# -------- ctor validations --------
def test_encoder_layer_param_checks():
with pytest.raises(TypeError):
EncoderLayer("32", 4, 64, 0.1)
with pytest.raises(TypeError):
EncoderLayer(32, "4", 64, 0.1)
with pytest.raises(TypeError):
EncoderLayer(32, 4, "64", 0.1)
with pytest.raises(TypeError):
EncoderLayer(32, 4, 64, "0.1")
with pytest.raises(TypeError):
EncoderLayer(32, 4, 64, 0.1, layer_norm_style=123)
with pytest.raises(ValueError):
EncoderLayer(0, 4, 64, 0.1)
with pytest.raises(ValueError):
EncoderLayer(32, 0, 64, 0.1)
with pytest.raises(ValueError):
EncoderLayer(32, 4, 0, 0.1)
with pytest.raises(ValueError):
EncoderLayer(32, 4, 64, 1.0) # upper bound excluded
with pytest.raises(ValueError):
EncoderLayer(32, 4, 64, 0.1, layer_norm_style="middle")
def test_transformer_encoder_layers_count_checks():
with pytest.raises(TypeError):
TransformerEncoder(32, 4, 64, "2", 0.1)
with pytest.raises(ValueError):
TransformerEncoder(32, 4, 64, 0, 0.1)
# -------- forward path --------
@pytest.mark.parametrize("B,S,D,H,FF,L", [(2, 5, 24, 3, 48, 2)])
def test_encoder_forward_happy_path(B, S, D, H, FF, L):
device = _dev()
enc = TransformerEncoder(D, H, FF, L, 0.1).to(device)
x = torch.randn(B, S, D, device=device)
heads = enc.layers[0].attention_layer.num_heads
src_pad = torch.zeros(B, heads, 1, S, dtype=torch.bool, device=device)
out = enc(x, src_pad)
assert out.shape == (B, S, D)
assert out.device == device
def test_transformer_encoder_pre_norm_layers_flag():
enc = TransformerEncoder(24, 3, 48, 2, 0.1, layer_norm_style="pre")
assert all(layer.pre_norm for layer in enc.layers)
def test_encoder_pre_norm_forward_matches_shapes():
device = _dev()
layer = EncoderLayer(24, 3, 48, 0.1, layer_norm_style="pre").to(device)
x = torch.randn(2, 5, 24, device=device)
mask = torch.zeros(2, 3, 1, 5, dtype=torch.bool, device=device)
out = layer(x, mask)
assert out.shape == x.shape
assert layer.pre_norm is True
def test_encoder_forward_input_checks_and_message_format():
layer = EncoderLayer(24, 3, 48, 0.1)
with pytest.raises(TypeError):
layer("not a tensor", None)
with pytest.raises(ValueError) as ei:
layer(torch.randn(2, 3, 4, 5), None) # rank 4
assert "x must be a 3D torch.Tensor of shape (B, S, D)" in str(ei.value)
|