fix checkpoint generated by ddp training cannot be loaded (#1682)

* fix checkpoint generated by ddp training cannot be loaded

* fix typo

* reformatted nerfstudio/pipelines/base_pipeline.py

* rewritten as for loop

* add test case

* reformatted

* fix pylint test failed

nerfstudio/pipelines/base_pipeline.py

@@ -108,9 +108,22 @@ class Pipeline(nn.Module):
         return self.model.device
 
     def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True):
-        model_state = {key[len("_model.") :]: value for key, value in state_dict.items() if key.startswith("_model.")}
+        is_ddp_model_state = True
+        model_state = {}
+        for key, value in state_dict.items():
+            if key.startswith("_model."):
+                # remove the "_model." prefix from key
+                model_state[key[len("_model.") :]] = value
+                # make sure that the "module." prefix comes from DDP,
+                # rather than an attribute of the model named "module"
+                if not key.startswith("_model.module."):
+                    is_ddp_model_state = False
+        # remove "module." prefix added by DDP
+        if is_ddp_model_state:
+            model_state = {key[len("module.") :]: value for key, value in model_state.items()}
+
         pipeline_state = {key: value for key, value in state_dict.items() if not key.startswith("_model.")}
-        self._model.load_state_dict(model_state, strict=strict)
+        self.model.load_state_dict(model_state, strict=strict)
         super().load_state_dict(pipeline_state, strict=False)
 
     @profiler.time_function
@@ -373,7 +386,7 @@ class VanillaPipeline(Pipeline):
         state = {
             (key[len("module.") :] if key.startswith("module.") else key): value for key, value in loaded_state.items()
         }
-        self._model.update_to_step(step)
+        self.model.update_to_step(step)
         self.load_state_dict(state, strict=True)
 
     def get_training_callbacks(

tests/pipelines/test_vanilla_pipeline.py

@@ -57,6 +57,9 @@ def test_load_state_dict():
         def __init__(self, *args, **kwargs):
             super().__init__(*args, **kwargs)
             self.register_parameter("param", nn.Parameter(torch.ones((3,))))
+            self.register_module(
+                "module", Model(*args, **kwargs)
+            )  # make sure that non-DDP checkpoint won't be processed as DDP one
 
         def load_state_dict(self, state_dict, strict=True):
             """Mocked load_state_dict"""
@@ -64,6 +67,7 @@ def test_load_state_dict():
             was_called = True
             assert strict
             assert "param" in state_dict
+            assert "module.device_indicator_param" in state_dict
 
     config = VanillaPipelineConfig(
         datamanager=VanillaDataManagerConfig(
@@ -77,3 +81,19 @@ def test_load_state_dict():
     pipeline.load_pipeline(state_dict, 0)
     assert was_called
     assert pipeline.model.param[0].item() == 2
+
+    # preparation for another test
+    state_dict["_model.param"].mul_(2)  # pylint: disable=unsubscriptable-object
+    was_called = False
+    # pretends to be a DDP checkpoint
+    ddp_state_dict = {}
+    for key, value in state_dict.items():
+        if key.startswith("_model."):
+            # replace the prefix "_model" with "_model.module"
+            ddp_state_dict["_model.module." + key[len("_model.") :]] = value
+        else:
+            ddp_state_dict[key] = value
+    # load DDP checkpoint
+    pipeline.load_pipeline(ddp_state_dict, 0)
+    assert was_called
+    assert pipeline.model.param[0].item() == 4