diff --git a/nerfstudio/configs/experiment_config.py b/nerfstudio/configs/experiment_config.py index c3a400dd1..1fa3598f6 100644 --- a/nerfstudio/configs/experiment_config.py +++ b/nerfstudio/configs/experiment_config.py @@ -76,6 +76,8 @@ class ExperimentConfig(InstantiateConfig): """Alias for --pipeline.model.prompt""" relative_model_dir: Path = Path("nerfstudio_models/") """Relative path to save all checkpoints.""" + load_scheduler: bool = True + """Whether to load the scheduler state_dict to resume training, if exists""" def is_viewer_enabled(self) -> bool: """Checks if a viewer is enabled.""" diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index a89e4801c..489543f0a 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -406,7 +406,8 @@ class Trainer: # load the checkpoints for pipeline, optimizers, and gradient scalar self.pipeline.load_pipeline(loaded_state["pipeline"], loaded_state["step"]) self.optimizers.load_optimizers(loaded_state["optimizers"]) - self.optimizers.load_schedulers(loaded_state["schedulers"]) + if "schedulers" in loaded_state and self.config.load_scheduler: + self.optimizers.load_schedulers(loaded_state["schedulers"]) self.grad_scaler.load_state_dict(loaded_state["scalers"]) CONSOLE.print(f"Done loading Nerfstudio checkpoint from {load_path}") elif load_checkpoint is not None: @@ -416,6 +417,8 @@ class Trainer: # load the checkpoints for pipeline, optimizers, and gradient scalar self.pipeline.load_pipeline(loaded_state["pipeline"], loaded_state["step"]) self.optimizers.load_optimizers(loaded_state["optimizers"]) + if "schedulers" in loaded_state and self.config.load_scheduler: + self.optimizers.load_schedulers(loaded_state["schedulers"]) self.grad_scaler.load_state_dict(loaded_state["scalers"]) CONSOLE.print(f"Done loading Nerfstudio checkpoint from {load_checkpoint}") else: