diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 7b35d98cc..c1cac0369 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -106,14 +106,12 @@ class Pipeline(nn.Module): return self.model.device def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True): - if "model" in state_dict: - tmp_state_dict = dict(state_dict) - model_state = tmp_state_dict.pop("model") - self.model.load_state_dict(model_state, strict=strict) - - # Model will be missing from the dict; therefore strict=False here. - strict = False - super().load_state_dict(state_dict, strict=strict) + model_state = { + key.replace("_model.", ""): value for key, value in state_dict.items() if key.startswith("_model.") + } + 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) + super().load_state_dict(pipeline_state, strict=False) @profiler.time_function def get_train_loss_dict(self, step: int): @@ -372,10 +370,7 @@ class VanillaPipeline(Pipeline): """ state = {key.replace("module.", ""): value for key, value in loaded_state.items()} self._model.update_to_step(step) - model_state = state.pop("model", {}) - self.model.load_state_dict(model_state, strict=True) - self.load_state_dict(state, strict=False) - state["model"] = model_state + self.load_state_dict(state, strict=True) def get_training_callbacks( self, training_callback_attributes: TrainingCallbackAttributes diff --git a/tests/pipelines/test_vanilla_pipeline.py b/tests/pipelines/test_vanilla_pipeline.py index 64168c300..69581ce27 100644 --- a/tests/pipelines/test_vanilla_pipeline.py +++ b/tests/pipelines/test_vanilla_pipeline.py @@ -9,6 +9,7 @@ from pathlib import Path # pylint: disable=unused-argument # pylint: disable=abstract-method import torch +from torch import nn from nerfstudio.cameras.cameras import Cameras from nerfstudio.data.datasets.base_dataset import DataparserOutputs, InputDataset @@ -53,11 +54,16 @@ def test_load_state_dict(): class MockedModel(Model): # """Mocked model""" - def load_state_dict(self, *args, strict=True): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.register_parameter("param", nn.Parameter(torch.ones((3,)))) + + def load_state_dict(self, state_dict, strict=True): """Mocked load_state_dict""" nonlocal was_called was_called = True assert strict + assert "param" in state_dict config = VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( @@ -67,5 +73,7 @@ def test_load_state_dict(): ) pipeline = VanillaPipeline(config, "cpu") state_dict = pipeline.state_dict() + state_dict["_model.param"].mul_(2) # pylint: disable=unsubscriptable-object pipeline.load_pipeline(state_dict, 0) assert was_called + assert pipeline.model.param[0].item() == 2