From 365e4c6663dac0c9c1dc19b1e54a386c0eafa8d0 Mon Sep 17 00:00:00 2001 From: Jeff Cook Date: Thu, 13 Jul 2023 16:22:36 -0600 Subject: [PATCH] Try non-strict model loading on RuntimeError. ... (#2214) * Try non-strict model loading on RuntimeError. ... This allows tcnn models that contained a torch hash_table to load. In particular, it allows us to load tcnn models prior to PR #2175. cf. https://github.com/nerfstudio-project/nerfstudio/pull/2175 * Only allow non-strict loading when called with defaults. * Reflow to avoid code duplication; blacken. --------- Co-authored-by: Matthew Tancik --- nerfstudio/pipelines/base_pipeline.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 363c8c9ea..11ba841e0 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -106,7 +106,7 @@ 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): + def load_state_dict(self, state_dict: Mapping[str, Any], strict: Optional[bool] = None): is_ddp_model_state = True model_state = {} for key, value in state_dict.items(): @@ -122,7 +122,15 @@ class Pipeline(nn.Module): 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) + + try: + self.model.load_state_dict(model_state, strict=True) + except RuntimeError: + if not strict: + self.model.load_state_dict(model_state, strict=False) + else: + raise + super().load_state_dict(pipeline_state, strict=False) @profiler.time_function @@ -410,7 +418,7 @@ class VanillaPipeline(Pipeline): (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) + self.load_state_dict(state) def get_training_callbacks( self, training_callback_attributes: TrainingCallbackAttributes