tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
Code

Load schedulers from ckpt in a backward-compatible manner (#2236)

Browse Source
This commit is contained in:
hyperzy 2023-07-19 02:02:15 +08:00 committed by GitHub
parent 9f487a464f
commit e7fd1aefa2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
nerfstudio/configs/experiment_config.py
View File

@ -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."""

5
nerfstudio/engine/trainer.py
View File

@ -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: