From e7fd1aefa29abf3c66c503934af49cd0aaa91c90 Mon Sep 17 00:00:00 2001 From: hyperzy <50564821+hyperzy@users.noreply.github.com> Date: Wed, 19 Jul 2023 02:02:15 +0800 Subject: [PATCH] Load schedulers from ckpt in a backward-compatible manner (#2236) --- nerfstudio/configs/experiment_config.py | 2 ++ nerfstudio/engine/trainer.py | 5 ++++- 2 files changed, 6 insertions(+), 1 deletion(-) 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: