diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 1da017ebc..c572bebeb 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/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( diff --git a/tests/pipelines/test_vanilla_pipeline.py b/tests/pipelines/test_vanilla_pipeline.py index 69581ce27..c51038135 100644 --- a/tests/pipelines/test_vanilla_pipeline.py +++ b/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