Fix encoders (mobilenet, inceptionv4)

Author: qubvel

segmentation_models_pytorch/encoders/inceptionv4.py

@@ -35,7 +35,6 @@
 class InceptionV4Encoder(InceptionV4, EncoderMixin):
     def __init__(
         self,
-        out_indexes: List[int],
         out_channels: List[int],
         depth: int = 5,
         output_stride: int = 32,
@@ -47,7 +46,7 @@ def __init__(
         self._in_channels = 3
         self._out_channels = out_channels
         self._output_stride = output_stride
-        self._out_indexes = out_indexes
+        self._out_indexes = [2, 4, 8, 14, len(self.features) - 1]
 
         # correct paddings
         for m in self.modules():
@@ -115,7 +114,6 @@ def load_state_dict(self, state_dict, **kwargs):
             },
         },
         "params": {
-            "out_indexes": [2, 4, 8, 14],
             "out_channels": [3, 64, 192, 384, 1024, 1536],
             "num_classes": 1001,
         },

segmentation_models_pytorch/encoders/mobilenet.py

@@ -40,7 +40,7 @@ def __init__(
         self._in_channels = 3
         self._out_channels = out_channels
         self._output_stride = output_stride
-        self._out_indexes = [2, 4, 7, 14]
+        self._out_indexes = [1, 3, 6, 13, len(self.features) - 1]
 
         del self.classifier
 

tests/models/base.py

@@ -148,14 +148,17 @@ def test_classification_head(self):
 
     def test_any_resolution(self):
         model = self.get_default_model()
-        if model.requires_divisible_input_shape:
-            self.skipTest("Model requires divisible input shape")
 
         sample = self._get_sample(
             height=self.default_height + 3,
             width=self.default_width + 7,
         ).to(default_device)
 
+        if model.requires_divisible_input_shape:
+            with self.assertRaises(RuntimeError, msg="Wrong input shape"):
+                output = model(sample)
+            return
+
         with torch.inference_mode():
             output = model(sample)