Update get_stages

Author: qubvel

segmentation_models_pytorch/encoders/_base.py

@@ -1,3 +1,6 @@
+import torch
+from typing import Sequence
+
 from . import _utils as utils
 
 
@@ -31,28 +34,21 @@ def set_in_channels(self, in_channels, pretrained=True):
             model=self, new_in_channels=in_channels, pretrained=pretrained
         )
 
-    def get_stages(self):
-        """Override it in your implementation"""
+    def get_stages(self) -> dict[int, Sequence[torch.nn.Module]]:
+        """Override it in your implementation, should return a dictionary with keys as
+        the output stride and values as the list of modules
+        """
         raise NotImplementedError
 
     def make_dilated(self, output_stride):
-        if output_stride == 16:
-            stage_list = [5]
-            dilation_list = [2]
-
-        elif output_stride == 8:
-            stage_list = [4, 5]
-            dilation_list = [2, 4]
-
-        else:
-            raise ValueError(
-                "Output stride should be 16 or 8, got {}.".format(output_stride)
-            )
-
-        self._output_stride = output_stride
+        if output_stride not in [8, 16]:
+            raise ValueError(f"Output stride should be 16 or 8, got {output_stride}.")
 
         stages = self.get_stages()
-        for stage_indx, dilation_rate in zip(stage_list, dilation_list):
-            utils.replace_strides_with_dilation(
-                module=stages[stage_indx], dilation_rate=dilation_rate
-            )
+        for stage_stride, stage_modules in stages.items():
+            if stage_stride <= output_stride:
+                continue
+
+            dilation_rate = stage_stride // output_stride
+            for module in stage_modules:
+                utils.replace_strides_with_dilation(module, dilation_rate)

segmentation_models_pytorch/encoders/dpn.py

@@ -43,6 +43,12 @@ def __init__(self, stage_idxs, out_channels, depth=5, **kwargs):
 
         del self.last_linear
 
+    def get_stages(self):
+        return {
+            16: self.features[self._stage_idxs[1] : self._stage_idxs[2]],
+            32: self.features[self._stage_idxs[2] : self._stage_idxs[3]],
+        }
+
     def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
         features = [x]
 

segmentation_models_pytorch/encoders/efficientnet.py

@@ -44,6 +44,12 @@ def __init__(self, stage_idxs, out_channels, model_name, depth=5):
 
         del self._fc
 
+    def get_stages(self):
+        return {
+            16: self._blocks[self._stage_idxs[1] : self._stage_idxs[2]],
+            32: self._blocks[self._stage_idxs[2] :],
+        }
+
     def apply_blocks(
         self, x: torch.Tensor, start_idx: int, end_idx: int
     ) -> torch.Tensor:

segmentation_models_pytorch/encoders/mix_transformer.py

@@ -526,6 +526,12 @@ def __init__(self, out_channels, depth=5, **kwargs):
         self._depth = depth
         self._in_channels = 3
 
+    def get_stages(self):
+        return {
+            16: [self.patch_embed3, self.block3, self.norm3],
+            32: [self.patch_embed4, self.block4, self.norm4],
+        }
+
     def forward(self, x):
         # create dummy output for the first block
         batch_size, _, height, width = x.shape

segmentation_models_pytorch/encoders/mobilenet.py

@@ -36,6 +36,12 @@ def __init__(self, out_channels, depth=5, **kwargs):
         self._in_channels = 3
         del self.classifier
 
+    def get_stages(self):
+        return {
+            16: self.features[7:14],
+            32: self.features[14:],
+        }
+
     def forward(self, x):
         features = [x]
 

segmentation_models_pytorch/encoders/mobileone.py

@@ -355,6 +355,12 @@ def __init__(
             num_se_blocks=num_blocks_per_stage[3] if use_se else 0,
         )
 
+    def get_stages(self):
+        return {
+            16: self.stage3,
+            32: self.stage4,
+        }
+
     def _make_stage(
         self, planes: int, num_blocks: int, num_se_blocks: int
     ) -> nn.Sequential:

segmentation_models_pytorch/encoders/resnet.py

@@ -44,6 +44,12 @@ def __init__(self, out_channels, depth=5, **kwargs):
         del self.fc
         del self.avgpool
 
+    def get_stages(self):
+        return {
+            16: self.layer3,
+            32: self.layer4,
+        }
+
     def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
         features = [x]
 

segmentation_models_pytorch/encoders/senet.py

@@ -23,8 +23,6 @@
         depth = 3 -> number of feature tensors = 4 (one with same resolution as input and 3 downsampled).
 """
 
-import torch.nn as nn
-
 from pretrainedmodels.models.senet import (
     SENet,
     SEBottleneck,
@@ -46,14 +44,10 @@ def __init__(self, out_channels, depth=5, **kwargs):
         del self.avg_pool
 
     def get_stages(self):
-        return [
-            nn.Identity(),
-            self.layer0[:-1],
-            nn.Sequential(self.layer0[-1], self.layer1),
-            self.layer2,
-            self.layer3,
-            self.layer4,
-        ]
+        return {
+            16: self.layer3,
+            32: self.layer4,
+        }
 
     def forward(self, x):
         features = [x]

segmentation_models_pytorch/encoders/timm_efficientnet.py

@@ -105,6 +105,12 @@ def __init__(self, stage_idxs, out_channels, depth=5, **kwargs):
 
         del self.classifier
 
+    def get_stages(self):
+        return {
+            16: self.blocks[self._stage_idxs[1] : self._stage_idxs[2]],
+            32: self.blocks[self._stage_idxs[2] :],
+        }
+
     def forward(self, x):
         features = [x]
 

segmentation_models_pytorch/encoders/timm_sknet.py

@@ -13,6 +13,12 @@ def __init__(self, out_channels, depth=5, **kwargs):
         del self.fc
         del self.global_pool
 
+    def get_stages(self):
+        return {
+            16: self.layer3,
+            32: self.layer4,
+        }
+
     def forward(self, x):
         features = [x]