test: add decoder and transformer tests

Author: okunator

cellseg_models_pytorch/decoders/tests/test_decoders.py

@@ -6,9 +6,11 @@
 
 @pytest.mark.parametrize("long_skip", ["unet", "unetpp", "unet3p", "unet3p-lite"])
 @pytest.mark.parametrize("merge_policy", ["cat", "sum"])
-def test_decoder_fwdbwd(long_skip, merge_policy):
-    enc_channels = (64, 32, 16, 8, 8)
-    out_dims = [256 // 2**i for i in range(6)][::-1]
+@pytest.mark.parametrize("use_conv", [True, False])
+@pytest.mark.parametrize("use_tr", [True, False])
+def test_decoder_fwdbwd(long_skip, merge_policy, use_conv, use_tr):
+    enc_channels = (64, 32, 64, 32)
+    out_dims = [32 // 2**i for i in range(4)][::-1]
 
     decoder1_kwargs = {"merge_policy": merge_policy}
     decoder2_kwargs = {"merge_policy": merge_policy}
@@ -23,17 +25,30 @@ def test_decoder_fwdbwd(long_skip, merge_policy):
         decoder5_kwargs,
     )
 
+    n_layers = None
+    n_blocks = None
+    if use_conv:
+        n_layers = (1, 1, 1, 1)
+        n_blocks = ((2,), (2,), (2,), (2,))
+
+    n_tr_layers = None
+    n_tr_blocks = None
+    if use_tr:
+        n_tr_layers = (1, 1, 1, 1)
+        n_tr_blocks = ((1,), (1,), (1,), (1,))
+
     decoder = Decoder(
         enc_channels=enc_channels,
-        model_input_size=256,
-        out_channels=(64, 32, 16, 8, 8),
-        n_layers=(1, 1, 1, 1, 1),
-        n_blocks=((2,), (2,), (2,), (2,), (2,)),
+        out_channels=(32, 32, 32, 32),
+        n_layers=n_layers,
+        n_blocks=n_blocks,
+        n_transformers=n_tr_layers,
+        n_transformer_blocks=n_tr_blocks,
         long_skip=long_skip,
         stage_params=stage_params,
     )
 
-    x = [torch.rand([1, enc_channels[i], out_dims[i], out_dims[i]]) for i in range(5)]
+    x = [torch.rand([1, enc_channels[i], out_dims[i], out_dims[i]]) for i in range(4)]
     out = decoder(*x)
 
     out[-1].mean().backward()

cellseg_models_pytorch/modules/tests/test_transformer.py

@@ -0,0 +1,27 @@
+import pytest
+import torch
+
+from cellseg_models_pytorch.modules import Transformer2D
+
+
+@pytest.mark.parametrize("block_type", ["basic", "slice"])
+def test_transformer(block_type):
+    in_channels = 64
+    B = 4
+    H = W = 32
+
+    x = torch.rand([B, in_channels, H, W])
+    tr = Transformer2D(
+        in_channels=in_channels,
+        num_heads=4,
+        head_dim=32,
+        n_blocks=1,
+        block_types=(block_type,),
+        biases=(False,),
+        dropouts=(0.0,),
+        slice_size=4,
+    )
+
+    out = tr(x)
+
+    assert out.shape == x.shape