Fixes bug introduced in #1499 preventing load_from_checkpoint to work correctly (#1501)

* Call load_state_dict on model classes

* Fixing bug introduced in #1499

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

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