test: add some unittests

Author: okunator

cellseg_models_pytorch/modules/conv_base.py

@@ -31,7 +31,7 @@ def __init__(
         bias: bool = False,
         attention: str = None,
         preattend: bool = False,
-        **kwargs
+        **kwargs,
     ) -> None:
         """Conv-block (basic) parent class.
 
@@ -143,7 +143,7 @@ def __init__(
         bias: bool = False,
         attention: str = None,
         preattend: bool = False,
-        **kwargs
+        **kwargs,
     ) -> None:
         """Bottleneck conv block parent-class.
 
@@ -305,7 +305,7 @@ def __init__(
         kernel_size: int = 3,
         attention: str = None,
         preattend: bool = False,
-        **kwargs
+        **kwargs,
     ) -> None:
         """Depthwise separable conv block parent class.
 
@@ -427,7 +427,7 @@ def __init__(
         kernel_size: int = 3,
         attention: str = None,
         preattend: bool = False,
-        **kwargs
+        **kwargs,
     ) -> None:
         """Mobile inverted bottleneck conv parent-class.
 
@@ -578,7 +578,7 @@ def __init__(
         activation: str = "relu",
         attention: str = None,
         preattend: bool = False,
-        **kwargs
+        **kwargs,
     ) -> None:
         """Fused mobile inverted conv block parent-class.
 
@@ -714,7 +714,7 @@ def __init__(
         kernel_size: int = 3,
         attention: str = None,
         preattend: bool = False,
-        **kwargs
+        **kwargs,
     ) -> None:
         """Dense block of the HoVer-Net.
 
@@ -830,135 +830,3 @@ def forward_features_preact(self, x: torch.Tensor) -> torch.Tensor:
         x = self.conv(x)
 
         return x
-
-
-class BasicConvOld(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        same_padding: bool = True,
-        normalization: str = "bn",
-        activation: str = "relu",
-        convolution: str = "conv",
-        preactivate: bool = False,
-        kernel_size=3,
-        groups: int = 1,
-        bias: bool = False,
-        attention: str = None,
-        preattend: bool = False,
-        **kwargs
-    ) -> None:
-        """Conv-block (basic) parent class.
-
-        Parameters
-        ----------
-            in_channels : int
-                Number of input channels.
-            out_channels : int
-                Number of output channels.
-            same_padding : bool, default=True
-                if True, performs same-covolution.
-            normalization : str, default="bn":
-                Normalization method.
-                One of: "bn", "bcn", "gn", "in", "ln", None
-            activation : str, default="relu"
-                Activation method.
-                One of: "mish", "swish", "relu", "relu6", "rrelu", "selu",
-                "celu", "gelu", "glu", "tanh", "sigmoid", "silu", "prelu",
-                "leaky-relu", "elu", "hardshrink", "tanhshrink", "hardsigmoid"
-            convolution : str, default="conv"
-                The convolution method. One of: "conv", "wsconv", "scaled_wsconv"
-            preactivate : bool, default=False
-                If True, normalization will be applied before convolution.
-            kernel_size : int, default=3
-                The size of the convolution kernel.
-            groups : int, default=1
-                Number of groups the kernels are divided into. If `groups == 1`
-                normal convolution is applied. If `groups = in_channels`
-                depthwise convolution is applied.
-            bias : bool, default=False,
-                Include bias term in the convolution.
-            attention : str, default=None
-                Attention method. One of: "se", "scse", "gc", "eca", "msca", None
-            preattend : bool, default=False
-                If True, Attention is applied at the beginning of forward pass.
-        """
-        super().__init__()
-        self.conv_choice = convolution
-        self.out_channels = out_channels
-        self.preattend = preattend
-        self.preactivate = preactivate
-
-        # set norm channel number for preactivation or normal
-        norm_channels = in_channels if preactivate else self.out_channels
-
-        # set padding. Works if dilation or stride are not adjusted
-        padding = (kernel_size - 1) // 2 if same_padding else 0
-
-        self.conv = Conv(
-            name=self.conv_choice,
-            in_channels=in_channels,
-            out_channels=out_channels,
-            kernel_size=kernel_size,
-            groups=groups,
-            padding=padding,
-            bias=bias,
-        )
-
-        self.norm = Norm(normalization, num_features=norm_channels)
-        self.act = Activation(activation)
-
-        # set attention channels
-        att_channels = in_channels if preattend else self.out_channels
-        self.att = Attention(attention, in_channels=att_channels)
-
-        self.downsample = None
-        if in_channels != out_channels:
-            self.downsample = nn.Sequential(
-                Conv(
-                    self.conv_choice,
-                    in_channels=in_channels,
-                    out_channels=out_channels,
-                    bias=False,
-                    kernel_size=1,
-                    padding=0,
-                ),
-                Norm(normalization, num_features=out_channels),
-            )
-
-    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
-        """Forward pass."""
-        identity = x
-        if self.downsample is not None:
-            identity = self.downsample(x)
-
-        x = self.att(x)
-
-        # residual
-        x = self.conv(x)
-        x = self.norm(x)
-
-        x += identity
-        x = self.act(x)
-
-        return x
-
-    def forward_features_preact(self, x: torch.Tensor) -> torch.Tensor:
-        """Forward pass with pre-activation."""
-        identity = x
-        if self.downsample is not None:
-            identity = self.downsample(x)
-
-        # pre-attention
-        x = self.att(x)
-
-        # preact residual
-        x = self.norm(x)
-        x = self.act(x)
-        x = self.conv(x)
-
-        x += identity
-        x = self.act(x)
-
-        return x

cellseg_models_pytorch/modules/tests/test_act.py

@@ -0,0 +1,82 @@
+import pytest
+import torch
+from torch.autograd import gradcheck
+from cellseg_models_pytorch.modules.act.swish import Swish
+from cellseg_models_pytorch.modules.act.mish import Mish
+from cellseg_models_pytorch.modules.act.gated_gelu import GEGLU
+
+@pytest.mark.parametrize("dim_in, dim_out", [
+    (64, 128),
+    (128, 256),
+])
+def test_geglu(dim_in, dim_out):
+    # Create the GEGLU layer
+    geglu = GEGLU(dim_in=dim_in, dim_out=dim_out)
+
+    # Create a random input tensor with the specified shape
+    x = torch.rand((1, 32, dim_in))
+
+    # Forward pass
+    output = geglu(x)
+
+    # Check the output shape
+    assert output.shape == (1, 32, dim_out)
+
+    # Check the output type
+    assert isinstance(output, torch.Tensor)
+
+
+@pytest.mark.parametrize("batch_size, num_features", [
+    (1, 10),
+    (2, 20),
+])
+def test_mish_fwdbwd(batch_size, num_features):
+    # Create the Mish layer
+    mish_layer = Mish()
+
+    # Create a random input tensor with the specified shape
+    x = torch.randn(batch_size, num_features, requires_grad=True)
+
+    # Forward pass
+    output = mish_layer(x)
+
+    # Check the output shape
+    assert output.shape == x.shape
+
+    # Check the output type
+    assert isinstance(output, torch.Tensor)
+
+    # Backward pass
+    output.sum().backward()
+    assert x.grad is not None
+
+    # Gradient check
+    # assert gradcheck(mish, (x,), eps=1e-6, atol=1e-4)
+
+
+@pytest.mark.parametrize("batch_size, num_features", [
+    (1, 10),
+    (2, 20),
+])
+def test_swish_fwdbwd(batch_size, num_features):
+    # Create the Swish layer
+    swish_layer = Swish()
+
+    # Create a random input tensor with the specified shape
+    x = torch.randn(batch_size, num_features, requires_grad=True)
+
+    # Forward pass
+    output = swish_layer(x)
+
+    # Check the output shape
+    assert output.shape == x.shape
+
+    # Check the output type
+    assert isinstance(output, torch.Tensor)
+
+    # Backward pass
+    output.sum().backward()
+    assert x.grad is not None
+
+    # Gradient check
+    # assert gradcheck(swish, (x,))

cellseg_models_pytorch/modules/tests/test_conv_blocks.py

@@ -2,6 +2,7 @@
 import torch
 
 from cellseg_models_pytorch.modules import ConvBlock
+from cellseg_models_pytorch.modules.mlp import ConvMlp
 
 
 @pytest.mark.parametrize(
@@ -43,3 +44,31 @@ def test_conv_block_fwdbwd(
 
     assert output.shape == torch.Size([1, out_channels, 16, 16])
     assert output.dtype == input.dtype
+
+
+@pytest.mark.parametrize("in_channels, out_channels, input_shape", [
+    (32, 16, (1, 32, 32, 32)),
+    (32, None, (2, 32, 32, 32)),
+])
+def test_convmlp(in_channels, out_channels, input_shape):
+    conv_mlp = ConvMlp(
+        in_channels=in_channels,
+        out_channels=out_channels,
+        mlp_ratio=1
+    )
+
+    if out_channels is None:
+        out_channels = in_channels
+
+    # Create a random input tensor with the specified shape
+    x = torch.rand(input_shape).float()
+
+    # Forward pass
+    output = conv_mlp(x)
+
+    # Check the output shape
+    expected_out_channels = in_channels if out_channels is None else out_channels
+    assert output.shape == (input_shape[0], expected_out_channels, input_shape[2], input_shape[3])
+
+    # Check the output type
+    assert isinstance(output, torch.Tensor)

cellseg_models_pytorch/modules/tests/test_miscmodules.py

@@ -7,7 +7,7 @@
     StyleBlock,
     StyleReshape,
 )
-
+from cellseg_models_pytorch.modules.patch_embeddings import PatchEmbed
 
 @pytest.mark.parametrize("in_channels", [32, 16])
 @pytest.mark.parametrize("out_channels", [16, 32])
@@ -37,3 +37,34 @@ def test_stylefwdbwd(in_channels, style_channels, out_channels):
     out.mean().backward()
 
     assert out.shape[1] == out_channels
+
+
+@pytest.mark.parametrize("in_channels, patch_size, head_dim, num_heads, input_shape", [
+    (3, 8, 32, 4, (1, 3, 128, 128)),
+    (3, 4, 16, 2, (1, 3, 64, 64)),
+])
+def test_patch_embed(in_channels, patch_size, head_dim, num_heads, input_shape):
+    # Create the PatchEmbed layer
+    patch_embed = PatchEmbed(
+        in_channels=in_channels,
+        patch_size=patch_size,
+        head_dim=head_dim,
+        num_heads=num_heads
+    )
+
+    # Create a random input tensor with the specified shape
+    x = torch.randn(input_shape)
+
+    # Forward pass
+    output = patch_embed(x)
+
+    # Calculate expected output shape
+    B, _, H, W = input_shape
+    expected_seq_len = (H // patch_size) * (W // patch_size)
+    expected_proj_dim = head_dim * num_heads
+
+    # Check the output shape
+    assert output.shape == (B, expected_seq_len, expected_proj_dim)
+
+    # Check the output type
+    assert isinstance(output, torch.Tensor)