fix(modules): adjust transformer layers

Author: okunator

cellseg_models_pytorch/decoders/decoder_stage.py

@@ -328,9 +328,9 @@ def forward(
                 Output torch.Tensor and extra skip torch.Tensors. If no extra
                 skips are present, returns None as the second return value.
         """
-        x = self.upsample(x)
+        x = self.upsample(x)  # (B, in_channels, H, W)
 
-        # long skip
+        # long skip (B, in_channels(+skip_channels), H, W)
         x = self.skip(x, ix=self.stage_ix, skips=skips, extra_skips=extra_skips)
 
         # unetpp returns extra skips
@@ -340,15 +340,15 @@ def forward(
         # conv layers
         if self.n_layers is not None:
             for conv_layer in self.conv_layers.values():
-                x = conv_layer(x, style)
+                x = conv_layer(x, style)  # (B, out_channels, H, W)
 
         # transformer layers
         if self.n_transformers is not None:
             for transformer in self.transformers.values():
-                x = transformer(x)
+                x = transformer(x)  # (B, long_skip_channels/out_channels, H, W)
 
         # channel pool if conv-layers are skipped.
         if self.n_layers is None:
-            x = self.ch_pool(x)
+            x = self.ch_pool(x)  # (B, out_channels, H, W)
 
         return x, extra_skips

cellseg_models_pytorch/modules/base_modules.py

@@ -1,12 +1,12 @@
 import torch
 import torch.nn as nn
 
-from .act import ACT_LOOKUP, TR_ACT_LOOKUP
+from .act import ACT_LOOKUP
 from .conv import CONV_LOOKUP
 from .norm import NORM_LOOKUP
 from .upsample import UP_LOOKUP
 
-__all__ = ["Activation", "Norm", "Up", "Conv", "Identity", "TransformerAct"]
+__all__ = ["Activation", "Norm", "Up", "Conv", "Identity"]
 
 
 class Identity(nn.Module):
@@ -151,32 +151,3 @@ def __init__(self, name: str, **kwargs) -> None:
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass for the convolution function."""
         return self.conv(x)
-
-
-class TransformerAct(nn.Module):
-    def __init__(self, name: str, **kwargs) -> None:
-        """Activation function for transformer outputs wrapper class.
-
-        Parameters:
-        -----------
-            name : str
-                Name of the transformer activation method.
-
-        Raises
-        ------
-            ValueError: if the transformer activation name is illegal.
-        """
-        super().__init__()
-
-        allowed = list(TR_ACT_LOOKUP.keys())
-        if name not in allowed:
-            raise ValueError(
-                "Illegal transformer activation method given. "
-                f"Allowed: {allowed}. Got: '{name}'"
-            )
-
-        self.tr_act = TR_ACT_LOOKUP[name](**kwargs)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Forward pass for the convolution function."""
-        return self.tr_act(x)

cellseg_models_pytorch/modules/norm/__init__.py

@@ -1,12 +1,13 @@
-from torch.nn import BatchNorm2d, InstanceNorm2d, SyncBatchNorm
+from torch.nn import BatchNorm2d, InstanceNorm2d, LayerNorm, SyncBatchNorm
 
 from .bcn import BCNorm
 from .gn import GroupNorm
 from .ln import LayerNorm2d
 
 NORM_LOOKUP = {
     "bn": BatchNorm2d,
-    "ln": LayerNorm2d,
+    "ln2d": LayerNorm2d,
+    "ln": LayerNorm,
     "bcn": BCNorm,
     "gn": GroupNorm,
     "in": InstanceNorm2d,
@@ -20,4 +21,5 @@
     "InstanceNorm2d",
     "SyncBatchNorm",
     "LayerNorm2d",
+    "LayerNorm",
 ]

cellseg_models_pytorch/modules/self_attention_modules.py

@@ -19,6 +19,7 @@ def __init__(
         dropout: float = 0.0,
         bias: bool = False,
         slice_size: int = 4,
+        **kwargs,
     ) -> None:
         """Compute self-attention.
 
@@ -52,6 +53,7 @@ def __init__(
                 `self_attention = "slice"`.
         """
         super().__init__()
+        self.out_channels = query_dim
         proj_channels = head_dim * num_heads
 
         # cross attention dim

cellseg_models_pytorch/modules/transformers.py

@@ -5,8 +5,8 @@
 
 from cellseg_models_pytorch.modules import SelfAttentionBlock
 
-from .base_modules import TransformerAct
-from .misc_modules import Proj2Attention
+from .mlp import MlpBlock
+from .patch_embeddings import ContiguousEmbed
 
 __all__ = ["Transformer2D", "TransformerLayer"]
 
@@ -25,7 +25,7 @@ def __init__(
         act: str = "geglu",
         num_groups: int = 32,
         slice_size: int = 4,
-        fc_projection_mult: int = 4,
+        mlp_ratio: int = 4,
         **kwargs,
     ) -> None:
         """Create a transformer for 2D-image-like (B, C, H, W) inputs.
@@ -69,15 +69,19 @@ def __init__(
                 layer.
         """
         super().__init__()
-        self.proj_in = Proj2Attention(
+        patch_norm = "gn" if in_channels >= 32 else None
+        self.patch_embed = ContiguousEmbed(
             in_channels=in_channels,
-            num_groups=num_groups,
+            patch_size=1,
             head_dim=head_dim,
             num_heads=num_heads,
+            normalization=patch_norm,
+            norm_kwargs={"num_features": in_channels, "num_groups": num_groups},
         )
+        self.proj_dim = self.patch_embed.proj_dim
 
         self.transformer = TransformerLayer(
-            query_dim=self.proj_in.proj_dim,
+            query_dim=self.proj_dim,
             num_heads=num_heads,
             head_dim=head_dim,
             cross_attention_dim=cross_attention_dim,
@@ -87,11 +91,11 @@ def __init__(
             biases=biases,
             act=act,
             slice_size=slice_size,
-            fc_projection_mult=fc_projection_mult,
+            mlp_ratio=mlp_ratio,
         )
 
         self.proj_out = nn.Conv2d(
-            self.proj_in.proj_dim, in_channels, kernel_size=1, stride=1, padding=0
+            self.proj_dim, in_channels, kernel_size=1, stride=1, padding=0
         )
 
     def forward(self, x: torch.Tensor, context: torch.Tensor = None) -> torch.Tensor:
@@ -114,13 +118,13 @@ def forward(self, x: torch.Tensor, context: torch.Tensor = None) -> torch.Tensor
         residual = x
 
         # 1. project
-        x = self.proj_in(x)
+        x = self.patch_embed(x)
 
         # 2. transformer
         x = self.transformer(x, context)
 
         # 3. Reshape back to image-like shape and project to original input channels.
-        x = x.reshape(B, H, W, self.proj_in.proj_dim).permute(0, 3, 1, 2)
+        x = x.reshape(B, H, W, self.proj_dim).permute(0, 3, 1, 2)
         x = self.proj_out(x)
 
         # 4. residual
@@ -134,13 +138,13 @@ def __init__(
         num_heads: int = 8,
         head_dim: int = 64,
         cross_attention_dim: int = None,
-        act: str = "geglu",
+        activation: str = "star_relu",
         n_blocks: int = 2,
         block_types: Tuple[str, ...] = ("basic", "basic"),
         dropouts: Tuple[float, ...] = (0.0, 0.0),
         biases: Tuple[bool, ...] = (False, False),
         slice_size: int = 4,
-        fc_projection_mult: int = 4,
+        mlp_ratio: int = 4,
         **kwargs,
     ) -> None:
         """Chain transformer blocks to compose a full generic transformer.
@@ -217,11 +221,12 @@ def __init__(
             )
             self.tr_blocks[f"transformer_{block_types[i]}_{i + 1}"] = att_block
 
-        proj_dim = int(query_dim * fc_projection_mult)
-        self.fc = nn.Sequential(
-            nn.LayerNorm(query_dim),
-            TransformerAct(act, dim_in=query_dim, dim_out=proj_dim),
-            nn.Linear(proj_dim, query_dim),
+        self.mlp = MlpBlock(
+            in_channels=query_dim,
+            mlp_ratio=mlp_ratio,
+            activation=activation,
+            normalization="ln",
+            norm_kwargs={"normalized_shape": query_dim},
         )
 
     def forward(self, x: torch.Tensor, context: torch.Tensor = None) -> torch.Tensor:
@@ -249,4 +254,4 @@ def forward(self, x: torch.Tensor, context: torch.Tensor = None) -> torch.Tensor
 
             x = tr_block(x, con)
 
-        return self.fc(x) + x
+        return self.mlp(x) + x