From 19e1d33132c5adecb75431e369e66f28843248c8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=C3=A1=C5=A1=20Kulh=C3=A1nek?= <jonas.kulhanek@live.com>
Date: Sat, 25 Feb 2023 12:23:06 +0100
Subject: [PATCH] Revert "Call load_state_dict on model classes (#1499)"

This reverts commit 93359b131e3b7f5b92264cb3520aa30419d1f1a9.
---
 nerfstudio/pipelines/base_pipeline.py    | 15 +----
 tests/pipelines/test_vanilla_pipeline.py | 71 ------------------------
 2 files changed, 1 insertion(+), 85 deletions(-)
 delete mode 100644 tests/pipelines/test_vanilla_pipeline.py

diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py
index 7b35d98cc..9d062e02f 100644
--- a/nerfstudio/pipelines/base_pipeline.py
+++ b/nerfstudio/pipelines/base_pipeline.py
@@ -21,7 +21,7 @@ import typing
 from abc import abstractmethod
 from dataclasses import dataclass, field
 from time import time
-from typing import Any, Dict, List, Mapping, Optional, Type, Union, cast
+from typing import Any, Dict, List, Optional, Type, Union, cast
 
 import torch
 import torch.distributed as dist
@@ -105,16 +105,6 @@ class Pipeline(nn.Module):
         """Returns the device that the model is on."""
         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)
-
     @profiler.time_function
     def get_train_loss_dict(self, step: int):
         """This function gets your training loss dict. This will be responsible for
@@ -372,10 +362,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
 
     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
deleted file mode 100644
index 64168c300..000000000
--- a/tests/pipelines/test_vanilla_pipeline.py
+++ /dev/null
@@ -1,71 +0,0 @@
-"""
-Test pipeline
-"""
-from pathlib import Path
-
-# pylint: disable=too-few-public-methods
-# pylint: disable=no-self-use
-# pylint: disable=missing-class-docstring
-# pylint: disable=unused-argument
-# pylint: disable=abstract-method
-import torch
-
-from nerfstudio.cameras.cameras import Cameras
-from nerfstudio.data.datasets.base_dataset import DataparserOutputs, InputDataset
-from nerfstudio.pipelines.base_pipeline import (
-    Model,
-    ModelConfig,
-    VanillaDataManagerConfig,
-    VanillaPipeline,
-    VanillaPipelineConfig,
-)
-
-
-class MockedDataManager:
-    """Mocked data manager"""
-
-    def __init__(self, *args, **kwargs):
-        num_images = 0
-        self.train_dataset = InputDataset(
-            DataparserOutputs(
-                image_filenames=[Path("filename.png")] * num_images,
-                cameras=Cameras(
-                    camera_to_worlds=torch.ones([num_images, 3, 4], dtype=torch.float32),
-                    fx=1.0,
-                    fy=1.0,
-                    cx=1.0,
-                    cy=1.0,
-                    width=2,
-                    height=2,
-                ),
-            )
-        )
-
-    def to(self, *args, **kwargs):
-        "Mocked to"
-        return self
-
-
-def test_load_state_dict():
-    """Test pipeline load_state_dict calls model's load_state_dict"""
-    was_called = False
-
-    class MockedModel(Model):  #
-        """Mocked model"""
-
-        def load_state_dict(self, *args, strict=True):
-            """Mocked load_state_dict"""
-            nonlocal was_called
-            was_called = True
-            assert strict
-
-    config = VanillaPipelineConfig(
-        datamanager=VanillaDataManagerConfig(
-            _target=MockedDataManager,
-        ),
-        model=ModelConfig(_target=MockedModel),
-    )
-    pipeline = VanillaPipeline(config, "cpu")
-    state_dict = pipeline.state_dict()
-    pipeline.load_pipeline(state_dict, 0)
-    assert was_called