IMP: filter pretrained models, cleanup

Author: JulienMaille

segmentation_models_pytorch/encoders/__init__.py

@@ -39,8 +39,11 @@
 
 def get_encoder(name, in_channels=3, depth=5, weights=None):
 
-    if name in timm_universal_encoders:
-        encoder = TimmUniversalEncoder(encoder_name=name, in_channels=in_channels, depth=depth, pretrained=weights is not None)
+    if name in timm_universal_encoders():
+        pretrained = weights is not None
+        if pretrained and name not in timm_universal_encoders(pretrained=True):
+            raise KeyError("No pretrained weights for encoder `{}`.".format(name))
+        encoder = TimmUniversalEncoder(encoder_name=name, in_channels=in_channels, depth=depth, pretrained=pretrained)
         global timm_setting
         timm_setting = encoder.formatted_settings
         return encoder
@@ -69,11 +72,11 @@ def get_encoder(name, in_channels=3, depth=5, weights=None):
 
 
 def get_encoder_names():
-    return list(encoders.keys()) + timm_universal_encoders
+    return list(encoders.keys()) + timm_universal_encoders()
 
 
 def get_preprocessing_params(encoder_name, pretrained="imagenet"):
-    if encoder_name in timm_universal_encoders:
+    if encoder_name in timm_universal_encoders():
         return timm_setting
 
     settings = encoders[encoder_name]["pretrained_settings"]

segmentation_models_pytorch/encoders/timm_universal.py

@@ -1,6 +1,7 @@
-from ._base import EncoderMixin
-from timm import create_model, list_models
 import torch.nn as nn
+from timm import create_model, list_models
+
+from ._base import EncoderMixin
 
 
 class TimmUniversalEncoder(nn.Module, EncoderMixin):
@@ -10,26 +11,27 @@ def __init__(self, encoder_name, in_channels, depth=5, pretrained=True, **kwargs
         self._in_channels = in_channels
 
         self.encoder = create_model(model_name=encoder_name[len(timm_universal_prefix):],
-                             in_chans=in_channels,
-                             exportable=True,   # onnx export
-                             features_only=True,
-                             pretrained=pretrained,
-                             out_indices=tuple(range(depth))) # FIXME need to handle a few special cases for specific models
+                                    in_chans=in_channels,
+                                    exportable=True,  # onnx export
+                                    features_only=True,
+                                    pretrained=pretrained,
+                                    out_indices=tuple(range(depth)) # FIXME need to handle a few special cases for specific models
+                                    )
 
         channels = self.encoder.feature_info.channels()
-        self._out_channels = (in_channels,)  + tuple(channels)
+        self._out_channels = (in_channels,) + tuple(channels)
 
         self.formatted_settings = {}
         self.formatted_settings["input_space"] = "RGB"
         self.formatted_settings["input_range"] = (0, 1)
         self.formatted_settings["mean"] = self.encoder.default_cfg['mean']
         self.formatted_settings["std"] = self.encoder.default_cfg['std']
 
-
     def forward(self, x):
         features = self.encoder(x)
         return [x] + features
 
 
 timm_universal_prefix = 'timm-u-'
-timm_universal_encoders = [f'{timm_universal_prefix}{i}' for i in list_models()]
+def timm_universal_encoders(**kwargs):
+    return [f'{timm_universal_prefix}{i}' for i in list_models(**kwargs)]

tests/test_models.py

@@ -18,13 +18,23 @@ def get_encoders():
         "resnext101_32x32d",
         "resnext101_32x48d",
     ]
-    encoders = smp.encoders.get_encoder_names()
+    encoders = list(smp.encoders.encoders.keys())
     if IS_TRAVIS:
         encoders = [e for e in encoders if e not in travis_exclude_encoders]
     return encoders
 
+def get_timm_u_encoders():
+    timm_exclude_encoders = [
+        'vit_*', 'tnt_*', 'pit_*',
+        '*iabn*', 'tresnet*',  # models using inplace abn
+        'dla*', 'hrnet*',  # hopefully fix at some point
+    ]
+
+    return smp.encoders.timm_universal_encoders(exclude_filters=timm_exclude_encoders)
+
 
 ENCODERS = get_encoders()
+ENCODERS_TIMM_U = get_timm_u_encoders()
 DEFAULT_ENCODER = "resnet18"
 
 
@@ -54,24 +64,26 @@ def _test_forward_backward(model, sample, test_shape=False):
         assert out.shape[2:] == sample.shape[2:]
 
 
-@pytest.mark.parametrize("encoder_name", ENCODERS)
+@pytest.mark.parametrize("encoder_name", ENCODERS + ENCODERS_TIMM_U)
 @pytest.mark.parametrize("encoder_depth", [3, 5])
 @pytest.mark.parametrize("model_class", [smp.FPN, smp.PSPNet, smp.Linknet, smp.Unet, smp.UnetPlusPlus])
 def test_forward(model_class, encoder_name, encoder_depth, **kwargs):
-    if model_class is smp.Unet or model_class is smp.UnetPlusPlus or model_class is smp.MAnet:
-        kwargs["decoder_channels"] = (16, 16, 16, 16, 16)[-encoder_depth:]
-    model = model_class(
-        encoder_name, encoder_depth=encoder_depth, encoder_weights=None, **kwargs
-    )
-    sample = get_sample(model_class)
-    model.eval()
-    if encoder_depth == 5 and model_class != smp.PSPNet:
-        test_shape = True
-    else:
-        test_shape = False
-
-    _test_forward(model, sample, test_shape)
-
+    try:
+        if model_class is smp.Unet or model_class is smp.UnetPlusPlus or model_class is smp.MAnet:
+            kwargs["decoder_channels"] = (16, 16, 16, 16, 16)[-encoder_depth:]
+        model = model_class(
+            encoder_name, encoder_depth=encoder_depth, encoder_weights=None, **kwargs
+        )
+        sample = get_sample(model_class)
+        model.eval()
+        if encoder_depth == 5 and model_class != smp.PSPNet:
+            test_shape = True
+        else:
+            test_shape = False
+        _test_forward(model, sample, test_shape)
+    except Exception as e:
+        print('\n\r{}-{}: Exception {}'.format(model_class.__name__, encoder_name, e))
+        assert False, 'Exception {}'.format(e)
 
 @pytest.mark.parametrize(
     "model_class",