Add BC for timm- encoders

Author: qubvel

segmentation_models_pytorch/base/model.py

@@ -90,3 +90,29 @@ def predict(self, x):
         x = self.forward(x)
 
         return x
+
+    def load_state_dict(self, state_dict, **kwargs):
+        # for compatibility of weights for
+        # timm- ported encoders with TimmUniversalEncoder
+        from segmentation_models_pytorch.encoders import TimmUniversalEncoder
+
+        if not isinstance(self.encoder, TimmUniversalEncoder):
+            return super().load_state_dict(state_dict, **kwargs)
+
+        patterns = ["regnet", "res2", "resnest", "mobilenetv3", "gernet"]
+
+        is_deprecated_encoder = any(
+            self.encoder.name.startswith(pattern) for pattern in patterns
+        )
+
+        if is_deprecated_encoder:
+            keys = list(state_dict.keys())
+            for key in keys:
+                new_key = key
+                if key.startswith("encoder.") and not key.startswith("encoder.model."):
+                    new_key = "encoder.model." + key.removeprefix("encoder.")
+                if "gernet" in self.encoder.name:
+                    new_key = new_key.replace(".stages.", ".stages_")
+                state_dict[new_key] = state_dict.pop(key)
+
+        return super().load_state_dict(state_dict, **kwargs)