style: minor docstring and code style fixes

Author: okunator

cellseg_models_pytorch/decoders/decoder.py

@@ -49,14 +49,14 @@ def __init__(
                 The number of convolution layers inside each of the decoder stages. The
                 argument can be given as a tuple, where each value indicates the number
                 of conv-layers inside each stage of the decoder allowing the mixing of
-                different sized layers inside the stages in the decoder. If set to None,
-                no conv-layers will be included in the decoder.
+                different sized layers inside the stages. If set to None, no conv-layers
+                will be included in the decoder.
             n_transformers : Union[None, int, Tuple[int, ...]] , optional
                 The number of transformer layers inside each of the decoder stages. The
                 argument can be given as a tuple, where each value indicates the number
-                of transformer-layers inside each stage of the decoder allowing the
-                mixing of different sized layers inside the stages in the decoder. If
-                set to None, no transformer layers will be included in the decoder.
+                of transformer-layers inside each stage of the decoder stages allowing
+                the mixing of different sized layers inside the stages. If set to None,
+                no transformer layers will be included in the decoder.
             n_conv_blocks : Union[int, Tuple[Tuple[int, ...], ...]], default=2
                 The number of blocks inside each conv-layer at each decoder stage. The
                 argument can be given as a nested tuple, where each value indicates the

cellseg_models_pytorch/decoders/decoder_stage.py

@@ -109,9 +109,10 @@ def __init__(
                 The tuple-length has to match `n_conv_layers`. Ignored if
                 `n_conv_layers` is None.
             attentions : Tuple[Tuple[str, ...], ...], default: ((None, "se"), )
-                Attention methods used inside the conv layers.
+                Channel-attention methods used inside the conv layers.
                 The tuple-length has to match `n_conv_layers`. Ignored if
-                `n_conv_layers` is None.
+                `n_conv_layers` is None. Allowed methods.: "se", "scse", "gc", "eca",
+                "msca", None.
             preactivates Tuple[Tuple[bool, ...], ...], default: ((False, False), )
                 Boolean flags for the conv layers to use pre-activation.
                 The tuple-length has to match `n_conv_layers`. Ignored if

cellseg_models_pytorch/modules/conv_base.py

@@ -64,7 +64,7 @@ def __init__(
             bias : bool, default=False,
                 Include bias term in the convolution.
             attention : str, default=None
-                Attention method. One of: "se", "scse", "gc", "eca", 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.
         """
@@ -186,7 +186,7 @@ def __init__(
             bias : bool, default=False,
                 Include bias term in the convolution.
             attention : str, default=None
-                Attention method. One of: "se", "scse", "gc", "eca", 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.
         """
@@ -336,7 +336,7 @@ def __init__(
             kernel_size : int, default=3
                 The size of the convolution kernel.
             attention : str, default=None
-                Attention method. One of: "se", "scse", "gc", "eca", 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.
         """
@@ -375,7 +375,7 @@ def __init__(
         self.act2 = Activation(activation)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
-        """Forward pass with pre-activation."""
+        """Forward pass."""
         if self.preattend:
             x = self.att(x)
 
@@ -394,7 +394,7 @@ def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
     def forward_features_preact(self, x: torch.Tensor) -> torch.Tensor:
-        """Forward pass."""
+        """Forward pass ith pre-activation."""
         if self.preattend:
             x = self.att(x)
 
@@ -459,7 +459,7 @@ def __init__(
             kernel_size : int, default=3
                 The size of the convolution kernel.
             attention : str, default=None
-                Attention method. One of: "se", "scse", "gc", "eca", 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.
         """
@@ -615,7 +615,7 @@ def __init__(
             preactivate : bool, default=False
                 If True, normalization will be applied before convolution.
             attention : str, default=None
-                Attention method. One of: "se", "scse", "gc", "eca", 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.
         """
@@ -750,7 +750,7 @@ def __init__(
             kernel_size : int, default=3
                 The size of the convolution kernel.
             attention : str, default=None
-                Attention method. One of: "se", "scse", "gc", "eca", 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.
         """
@@ -880,7 +880,7 @@ def __init__(
             bias : bool, default=False,
                 Include bias term in the convolution.
             attention : str, default=None
-                Attention method. One of: "se", "scse", "gc", "eca", 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.
         """

cellseg_models_pytorch/modules/conv_block.py

@@ -134,7 +134,7 @@ def __init__(
             bias : bool, default=True,
                 Include bias term in the convolution block. Only used for `BaasicConv`.
             attention : str, default=None
-                Attention method. One of: "se", "scse", "gc", "eca", 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.
             use_style : bool, default=False

cellseg_models_pytorch/modules/conv_layer.py

@@ -73,7 +73,7 @@ def __init__(
             biases : bool, default=(True, True)
                 Include bias terms in the convolution blocks.
             attentions : Tuple[str, ...], default=(None, None)
-                Attention method. One of: "se", "scse", "gc", "eca", None
+                Attention method. One of: "se", "scse", "gc", "eca", "msca", None
             preattends : Tuple[bool, ...], default=(False, False)
                 If True, Attention is applied at the beginning of forward pass.
             use_styles : bool, default=(False, False)

cellseg_models_pytorch/modules/misc_modules.py

@@ -28,13 +28,18 @@ def __init__(
                 Flag, whether the scaling is an inplace operation.
         """
         super().__init__()
+        self.dim = dim
         self.inplace = inplace
         self.gamma = nn.Parameter(init_values * torch.ones(dim))
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass of the layer scaling."""
         return x.mul_(self.gamma) if self.inplace else x * self.gamma
 
+    def extra_repr(self) -> str:
+        """Add extra to repr."""
+        return f"dim={self.dim}, inplace={self.inplace}"
+
 
 class ChannelPool(nn.Module):
     def __init__(

cellseg_models_pytorch/modules/patch_embeddings.py

@@ -59,7 +59,9 @@ def __init__(
             num_heads : int, default=8
                 Number of heads in multi-head self-attention.
             flatten : bool, default=True
-                If True, the output will be flattened to a sequence.
+                If True, the output will be flattened to a sequence. After flattening
+                output will have shape (B, H'*W', head_dim*num_heads). If False, the
+                output shape will remain (B, C, H', W').
             normalization : str, optional
                 The name of the normalization method.
                 One of: "bn", "bcn", "gn", "in", "ln", "lrn", None
@@ -105,6 +107,7 @@ def __init__(
         self.kernel_size = patch_size if kernel_size is None else kernel_size
         self.pad = pad
         self.stride = stride
+        norm_kwargs = norm_kwargs if norm_kwargs is not None else {}
 
         self.norm = Norm(normalization, **norm_kwargs)
         self.proj = nn.Conv2d(

cellseg_models_pytorch/modules/transformers.py

@@ -135,7 +135,7 @@ def forward(self, x: torch.Tensor, context: torch.Tensor = None) -> torch.Tensor
         B, _, H, W = x.shape
         residual = x
 
-        # 1. project
+        # 1. embed and project
         x = self.patch_embed(x)
 
         # 2. transformer
@@ -151,7 +151,7 @@ def forward(self, x: torch.Tensor, context: torch.Tensor = None) -> torch.Tensor
         # NOTE: the kernel_size, pad, & stride has to be set correctly for this to work
         if p_H < H:
             scale_factor = H // p_H
-            x = F.interpolate(x, scale_factor=scale_factor, mode="bilinear")
+            x = F.interpolate(x, scale_factor=int(scale_factor), mode="bilinear")
 
         # 4. project to original input channels
         x = self.proj_out(x)
@@ -268,15 +268,13 @@ def __init__(
             ls = LayerScale(query_dim) if layer_scales[i] else Identity()
             self.layer_scales.append(ls)
 
-            # self.tr_blocks[f"transformer_{block_types[i]}_{i + 1}"] = tr_block
-
         self.mlp = MlpBlock(
             in_channels=query_dim,
             mlp_ratio=mlp_ratio,
             activation=activation,
             normalization="ln",
             norm_kwargs={"normalized_shape": query_dim},
-            activation_kwargs=kwargs,
+            act_kwargs=kwargs,
         )
 
     def forward(self, x: torch.Tensor, context: torch.Tensor = None) -> torch.Tensor: