Repository: airctic/icevision

Issue:
load_state_dict bug
## 🐛 Bug
Introducing `load_state_dict` into the models in order to avoid doing `model.model.load_state_dict` caused a bug when trying to load from a previous checkpoint.

When saving the model the state dict looks like: `model.backbone...` but the new `load_state_dict` expects only `backbone`, without the model part.

More specifically, the following fails:
```python
model.load_state_dict(model.state_dict())
```

Related to #137 



Discussion:

Resolving patch:
diff --git a/mantisshrimp/models/mantis_module.py b/mantisshrimp/models/mantis_module.py
--- a/mantisshrimp/models/mantis_module.py
+++ b/mantisshrimp/models/mantis_module.py
@@ -6,13 +6,6 @@
 
 
 class MantisModule(DeviceModuleMixin, ParametersSplitsModuleMixin, nn.Module, ABC):
-    @abstractmethod
-    def load_state_dict(
-        self, state_dict: Dict[str, Tensor], strict: bool = True,
-    ):
-        """ Load weights
-        """
-
     @classmethod
     @abstractmethod
     def dataloader(cls, **kwargs) -> DataLoader:
diff --git a/mantisshrimp/models/mantis_rcnn/mantis_faster_rcnn.py b/mantisshrimp/models/mantis_rcnn/mantis_faster_rcnn.py
--- a/mantisshrimp/models/mantis_rcnn/mantis_faster_rcnn.py
+++ b/mantisshrimp/models/mantis_rcnn/mantis_faster_rcnn.py
@@ -54,11 +54,6 @@ def predict(self, images: List[np.ndarray], detection_threshold: float = 0.5):
             images=images, convert_raw_prediction=convert_raw_prediction
         )
 
-    def load_state_dict(
-        self, state_dict: Dict[str, Tensor], strict: bool = True,
-    ):
-        return self.model.load_state_dict(state_dict=state_dict, strict=strict)
-
     @property
     def param_groups(self):
         return self._param_groups
diff --git a/mantisshrimp/models/mantis_rcnn/mantis_mask_rcnn.py b/mantisshrimp/models/mantis_rcnn/mantis_mask_rcnn.py
--- a/mantisshrimp/models/mantis_rcnn/mantis_mask_rcnn.py
+++ b/mantisshrimp/models/mantis_rcnn/mantis_mask_rcnn.py
@@ -60,11 +60,6 @@ def predict(
             images=images, convert_raw_prediction=convert_raw_prediction
         )
 
-    def load_state_dict(
-        self, state_dict: Dict[str, Tensor], strict: bool = True,
-    ):
-        return self.model.load_state_dict(state_dict=state_dict, strict=strict)
-
     @property
     def param_groups(self):
         return self._param_groups


Test patch:
diff --git a/tests/models/mantis_rcnn/test_mantis_rcnn_predict.py b/tests/models/mantis_rcnn/test_mantis_rcnn_predict.py
--- a/tests/models/mantis_rcnn/test_mantis_rcnn_predict.py
+++ b/tests/models/mantis_rcnn/test_mantis_rcnn_predict.py
@@ -21,7 +21,7 @@ def pretrained_state_dict():
 
 def test_mantis_mask_rcnn_predict(sample_images, pretrained_state_dict):
     model = MantisMaskRCNN(91, min_size=128, max_size=128)
-    model.load_state_dict(pretrained_state_dict)
+    model.model.load_state_dict(pretrained_state_dict)
 
     preds = model.predict(sample_images)
 
diff --git a/tests/models/test_parameters_splits_module_mixin.py b/tests/models/test_parameters_splits_module_mixin.py
--- a/tests/models/test_parameters_splits_module_mixin.py
+++ b/tests/models/test_parameters_splits_module_mixin.py
@@ -18,11 +18,6 @@ def forward(self, x):
         def dataloader(cls, **kwargs) -> DataLoader:
             return DataLoader(**kwargs)
 
-        def load_state_dict(
-            self, state_dict: Dict[str, Tensor], strict: bool = True,
-        ):
-            pass
-
     return SimpleModel()