From 44ad94073d759fd6e6e755a6bb6343409e452377 Mon Sep 17 00:00:00 2001 From: Terrance Wang Date: Fri, 10 Mar 2023 23:09:46 -0800 Subject: [PATCH] bug fix for model loading from checkpoint (#1567) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * bug fix for model loading from checkpoint * removing prefix instead * fix in load_pipeline if ddp as well * fix * formatting --------- Co-authored-by: Jonáš Kulhánek --- nerfstudio/pipelines/base_pipeline.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index fd5d6fcc7..bb4a25f95 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -107,9 +107,7 @@ class Pipeline(nn.Module): return self.model.device def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True): - model_state = { - key.replace("_model.", ""): value for key, value in state_dict.items() if key.startswith("_model.") - } + model_state = {key[len("_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) @@ -370,7 +368,9 @@ class VanillaPipeline(Pipeline): loaded_state: pre-trained model state dict step: training step of the loaded checkpoint """ - state = {key.replace("module.", ""): value for key, value in loaded_state.items()} + 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.load_state_dict(state, strict=True)