refactor: verbose error msgs for abstract modules

Author: okunator

cellseg_models_pytorch/modules/attention_modules.py

@@ -325,7 +325,13 @@ def __init__(self, name: str, **kwargs) -> None:
             )
 
         if name is not None:
-            self.att = ATT_LOOKUP[name](**kwargs)
+            try:
+                self.att = ATT_LOOKUP[name](**kwargs)
+            except Exception as e:
+                raise Exception(
+                    "Encountered an error when trying to init chl attention function: "
+                    f"Attention(name='{name}'): {e.__class__.__name__}: {e}"
+                )
         else:
             self.att = Identity()
 

cellseg_models_pytorch/modules/base_modules.py

@@ -50,7 +50,13 @@ def __init__(self, name: str, **kwargs) -> None:
             try:
                 self.act = ACT_LOOKUP[name](**kwargs, inplace=True)
             except Exception:
-                self.act = ACT_LOOKUP[name](**kwargs)
+                try:
+                    self.act = ACT_LOOKUP[name](**kwargs)
+                except Exception as e:
+                    raise Exception(
+                        "Encountered an error when trying to init activation function: "
+                        f"Activation(name='{name}'): {e.__class__.__name__}: {e}"
+                    )
         else:
             self.act = Identity()
 
@@ -81,7 +87,13 @@ def __init__(self, name: str, **kwargs) -> None:
             )
 
         if name is not None:
-            self.norm = NORM_LOOKUP[name](**kwargs)
+            try:
+                self.norm = NORM_LOOKUP[name](**kwargs)
+            except Exception as e:
+                raise Exception(
+                    "Encountered an error when trying to init normalization function: "
+                    f"Norm(name='{name}'): {e.__class__.__name__}: {e}"
+                )
         else:
             self.norm = Identity()
 
@@ -118,7 +130,14 @@ def __init__(self, name: str, scale_factor: int = 2, **kwargs) -> None:
             kwargs["align_corners"] = True
 
         kwargs["scale_factor"] = scale_factor
-        self.up = UP_LOOKUP[name](**kwargs)
+
+        try:
+            self.up = UP_LOOKUP[name](**kwargs)
+        except Exception as e:
+            raise Exception(
+                "Encountered an error when trying to init upsampling function: "
+                f"Up(name='{name}'): {e.__class__.__name__}: {e}"
+            )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass for the upsampling function."""
@@ -146,7 +165,13 @@ def __init__(self, name: str, **kwargs) -> None:
                 f"Illegal convolution method given. Allowed: {allowed}. Got: '{name}'"
             )
 
-        self.conv = CONV_LOOKUP[name](**kwargs)
+        try:
+            self.conv = CONV_LOOKUP[name](**kwargs)
+        except Exception as e:
+            raise Exception(
+                "Encountered an error when trying to init convolution function: "
+                f"Conv(name='{name}'): {e.__class__.__name__}: {e}"
+            )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass for the convolution function."""

cellseg_models_pytorch/modules/conv_block.py

@@ -161,21 +161,27 @@ def __init__(
                 f"Illegal `short_skip` given. Got: '{short_skip}'. Allowed: {allowed}."
             )
 
-        self.block = CONVBLOCK_LOOKUP[name](
-            in_channels=in_channels,
-            out_channels=out_channels,
-            same_padding=same_padding,
-            normalization=normalization,
-            activation=activation,
-            convolution=convolution,
-            preactivate=preactivate,
-            kernel_size=kernel_size,
-            groups=groups,
-            bias=bias,
-            attention=attention,
-            preattend=preattend,
-            **kwargs,
-        )
+        try:
+            self.block = CONVBLOCK_LOOKUP[name](
+                in_channels=in_channels,
+                out_channels=out_channels,
+                same_padding=same_padding,
+                normalization=normalization,
+                activation=activation,
+                convolution=convolution,
+                preactivate=preactivate,
+                kernel_size=kernel_size,
+                groups=groups,
+                bias=bias,
+                attention=attention,
+                preattend=preattend,
+                **kwargs,
+            )
+        except Exception as e:
+            raise Exception(
+                "Encountered an error when trying to init ConvBlock module: "
+                f"ConvBlock(name='{name}'): {e.__class__.__name__}: {e}"
+            )
 
         self.downsample = None
         if short_skip == "residual" and in_channels != self.out_channels:

cellseg_models_pytorch/modules/self_attention/exact_attention.py

@@ -193,7 +193,13 @@ def forward(
                 The self-attention matrix. Same shape as inputs.
         """
         if self.self_attention == "memeff":
-            attn = memory_efficient_attention(query, key, value)
+            if all([query.is_cuda, key.is_cuda, value.is_cuda]):
+                attn = memory_efficient_attention(query, key, value)
+            else:
+                raise RuntimeError(
+                    "`xformers.ops.memory_efficient_attention` is only implemented "
+                    "for cuda. Make sure your inputs & model devices are set to cuda."
+                )
         elif self.self_attention == "flash":
             raise NotImplementedError
         elif self.self_attention == "slice":