Fix weight loading for deprecate encoders

Author: qubvel

segmentation_models_pytorch/base/model.py

@@ -3,7 +3,6 @@
 
 from . import initialization as init
 from .hub_mixin import SMPHubMixin
-from ..encoders.timm_universal import TimmUniversalEncoder
 
 T = TypeVar("T", bound="SegmentationModel")
 
@@ -82,18 +81,3 @@ 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
-        if isinstance(self.encoder, TimmUniversalEncoder):
-            keys = list(state_dict.keys())
-            for key in keys:
-                new_key = key
-                if key.startswith("encoder.") and not key.startswith("encoder.model."):
-                    new_key = key.replace("encoder.", "encoder.model.")
-                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)

segmentation_models_pytorch/encoders/timm_universal.py

@@ -196,6 +196,28 @@ def output_stride(self) -> int:
         """
         return min(self._output_stride, 2**self._depth)
 
+    def load_state_dict(self, state_dict, **kwargs):
+        # for compatibility of weights for
+        # timm- ported encoders with TimmUniversalEncoder
+
+        patterns = ["regnet", "res2", "resnest", "mobilenetv3", "gernet"]
+
+        is_deprecated_encoder = any(
+            self.name.startswith(pattern) for pattern in patterns
+        )
+
+        if is_deprecated_encoder:
+            keys = list(state_dict.keys())
+            for key in keys:
+                new_key = key
+                if not key.startswith("model."):
+                    new_key = "model." + key
+                if "gernet" in self.name:
+                    new_key = new_key.replace(".stages.", ".stages_")
+                state_dict[new_key] = state_dict.pop(key)
+
+        return super().load_state_dict(state_dict, **kwargs)
+
 
 def _merge_kwargs_no_duplicates(a: dict[str, Any], b: dict[str, Any]) -> dict[str, Any]:
     """