Fix silent fail if network loads wrong weights (#1521)

* Add warning if the imported network does not have the right keys Signed-off-by: Matthias Hadlich <[email protected]> * Update warning Signed-off-by: Matthias Hadlich <[email protected]> * Set load_strict to false for deepedit Signed-off-by: Matthias Hadlich <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Matthias Hadlich <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: tangy5 <[email protected]> Co-authored-by: SACHIDANAND ALLE <[email protected]>
Project-MONAI · Aug 26, 2023 · 7208641 · 7208641
1 parent 43e9e2a
commit 7208641
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/monailabel/tasks/infer/basic_infer.py b/monailabel/tasks/infer/basic_infer.py
@@ -57,7 +57,7 @@ def __init__(
         output_label_key: str = "pred",
         output_json_key: str = "result",
         config: Union[None, Dict[str, Any]] = None,
-        load_strict: bool = False,
+        load_strict: bool = True,
         roi_size=None,
         preload=False,
         train_mode=False,
@@ -453,6 +453,16 @@ def _get_network(self, device, data):
                 if path:
                     checkpoint = torch.load(path, map_location=torch.device(device))
                     model_state_dict = checkpoint.get(self.model_state_dict, checkpoint)
+
+                    if set(self.network.state_dict().keys()) != set(checkpoint.keys()):
+                        logger.warning(
+                            f"Checkpoint keys don't match network.state_dict()! Items that exist in only one dict"
+                            f" but not in the other: {set(self.network.state_dict().keys()) ^ set(checkpoint.keys())}"
+                        )
+                        logger.warning(
+                            "The run will now continue unless load_strict is set to True. "
+                            "If loading fails or the network behaves abnormally, please check the loaded weights"
+                        )
                     network.load_state_dict(model_state_dict, strict=self.load_strict)
             else:
                 network = torch.jit.load(path, map_location=torch.device(device))

diff --git a/sample-apps/radiology/lib/infers/deepedit.py b/sample-apps/radiology/lib/infers/deepedit.py
@@ -71,6 +71,7 @@ def __init__(
         self.spatial_size = spatial_size
         self.target_spacing = target_spacing
         self.number_intensity_ch = number_intensity_ch
+        self.load_strict = False
 
     def pre_transforms(self, data=None):
         t = [