feat(models): return featmaps at every dec stage

Author: okunator

cellseg_models_pytorch/decoders/decoder.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 import torch
 import torch.nn as nn
@@ -77,13 +77,14 @@ def __init__(
 
         self.out_channels = decoder_block.out_channels
 
-    def forward(
-        self, *features: Tuple[torch.Tensor], style: torch.Tensor = None
-    ) -> torch.Tensor:
-        """Forward pass of the decoder."""
+    def forward_features(
+        self, features: Tuple[torch.Tensor], style: torch.Tensor = None
+    ) -> List[torch.Tensor]:
+        """Forward pass of the decoder. Returns all the decoder stage feats."""
         head = features[0]
         skips = features[1:]
         extra_skips = [head] if self.long_skip == "unet3p" else []
+        ret_feats = []
 
         x = head
         for _, decoder_stage in enumerate(self.values()):
@@ -96,4 +97,14 @@ def forward(
             elif self.long_skip == "unet3p":
                 extra_skips.append(x)
 
-        return x
+            ret_feats.append(x)
+
+        return ret_feats
+
+    def forward(
+        self, *features: Tuple[torch.Tensor], style: torch.Tensor = None
+    ) -> torch.Tensor:
+        """Forward pass of the decoder."""
+        dec_feats = self.forward_features(features, style)
+
+        return dec_feats

cellseg_models_pytorch/decoders/tests/test_decoders.py

@@ -36,9 +36,9 @@ def test_decoder_fwdbwd(long_skip, merge_policy):
     x = [torch.rand([1, enc_channels[i], out_dims[i], out_dims[i]]) for i in range(5)]
     out = decoder(*x)
 
-    out.mean().backward()
+    out[-1].mean().backward()
 
-    assert out.shape[1] == decoder.out_channels
+    assert out[-1].shape[1] == decoder.out_channels
 
 
 @pytest.mark.slow
@@ -163,6 +163,6 @@ def test_decoder_fwdbwd_all(
     x = [torch.rand([1, enc_channels[i], out_dims[i], out_dims[i]]) for i in range(5)]
     out = decoder(*x)
 
-    out.mean().backward()
+    out[-1].mean().backward()
 
-    assert out.shape[1] == decoder.out_channels
+    assert out[-1].shape[1] == decoder.out_channels

cellseg_models_pytorch/models/base/_base_model.py

@@ -10,17 +10,38 @@
 
 
 class BaseMultiTaskSegModel(nn.ModuleDict):
+    def forward_encoder(self, x: torch.Tensor) -> List[torch.Tensor]:
+        """Forward the model encoder."""
+        self._check_input_shape(x)
+        feats = self.encoder(x)
+
+        return feats
+
+    def forward_style(self, feat: torch.Tensor) -> torch.Tensor:
+        """Forward the style domain adaptation layer.
+
+        NOTE: returns None if style channels are not given at model init.
+        """
+        style = None
+        if self.make_style is not None:
+            style = self.make_style(feat)
+
+        return style
+
     def forward_dec_features(
         self, feats: List[torch.Tensor], style: torch.Tensor = None
-    ) -> Dict[str, torch.Tensor]:
-        """Forward pass of the decoders in a multi-task seg model."""
+    ) -> Dict[str, List[torch.Tensor]]:
+        """Forward pass of all the decoder features mappings in the model.
+
+        NOTE: returns all the features from diff decoder stages in a list.
+        """
         res = {}
         decoders = [k for k in self.keys() if "decoder" in k]
 
         for dec in decoders:
-            x = self[dec](*feats, style=style)
+            featlist = self[dec](*feats, style=style)
             branch = dec.split("_")[0]
-            res[branch] = x
+            res[branch] = featlist
 
         return res
 
@@ -30,10 +51,9 @@ def forward_heads(
         """Forward pass of the seg heads in a multi-task seg model."""
         res = {}
         heads = [k for k in self.keys() if "head" in k]
-
         for head in heads:
             branch = head.split("_")[0]
-            x = self[head](dec_feats[branch])
+            x = self[head](dec_feats[branch][-1])  # the last decoder stage feat map
             res[branch] = x
 
         return res

cellseg_models_pytorch/models/base/_multitask_unet.py

@@ -125,14 +125,8 @@ def __init__(
 
     def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
         """Forward pass of Multi-task U-net."""
-        self._check_input_shape(x)
-
-        feats = self.encoder(x)
-
-        style = None
-        if self.make_style is not None:
-            style = self.make_style(feats[0])
-
+        feats = self.forward_encoder(x)
+        style = self.forward_style(feats[0])
         dec_feats = self.forward_dec_features(feats, style)
 
         for decoder_name in self.heads.keys():

cellseg_models_pytorch/models/cellpose/cellpose.py

@@ -201,14 +201,8 @@ def __init__(
 
     def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
         """Forward pass of Cellpose U-net."""
-        self._check_input_shape(x)
-
-        feats = self.encoder(x)
-
-        style = None
-        if self.make_style is not None:
-            style = self.make_style(feats[0])
-
+        feats = self.forward_encoder(x)
+        style = self.forward_style(feats[0])
         dec_feats = self.forward_dec_features(feats, style)
 
         for decoder_name in self.heads.keys():

cellseg_models_pytorch/models/hovernet/hovernet.py

@@ -194,13 +194,8 @@ def __init__(
 
     def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
         """Forward pass of HoVer-Net."""
-        self._check_input_shape(x)
-
-        feats = self.encoder(x)
-
-        style = None
-        if self.make_style is not None:
-            style = self.make_style(feats[0])
+        feats = self.forward_encoder(x)
+        style = self.forward_style(feats[0])
 
         dec_feats = self.forward_dec_features(feats, style)
         out = self.forward_heads(dec_feats)

cellseg_models_pytorch/models/stardist/stardist.py

@@ -209,20 +209,17 @@ def __init__(
 
     def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
         """Forward pass of Stardist."""
-        self._check_input_shape(x)
-
-        feats = self.encoder(x)
-
-        style = None
-        if self.make_style is not None:
-            style = self.make_style(feats[0])
+        feats = self.forward_encoder(x)
+        style = self.forward_style(feats[0])
 
         dec_feats = self.forward_dec_features(feats, style)
         # Extra convs after decoders
         for e in self.extra_convs.keys():
             for extra_conv in self.extra_convs[e].keys():
                 k = self.aux_key if extra_conv not in dec_feats.keys() else extra_conv
-                dec_feats[extra_conv] = self[f"{extra_conv}_features"](dec_feats[k])
+                dec_feats[extra_conv] = [
+                    self[f"{extra_conv}_features"](dec_feats[k][-1])
+                ]  # use last decoder feat
 
         # seg heads
         for decoder_name in self.heads.keys():

cellseg_models_pytorch/training/lit/lightning_experiment.py

@@ -74,7 +74,7 @@ def __init__(
                 Params dict for the scheduler. Refer to torch lr_scheduler docs
                 for the possible scheduler arguments.
             log_freq : int, default=100
-                Return soft masks every every n batches for callbacks and logging.
+                Return logs every n batches in logging callbacks.
 
         Raises
         ------

changelog.d/20221018_184607_oskari.lehtonen.md

@@ -0,0 +1,3 @@
+## Features
+
+- Support to return all of the feature maps from each decoder stage.