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

Add saving/loading scheduler state_dict (#2203)

Browse Source 
- Add saving/loading scheduler stat_dict
This commit is contained in:
hyperzy 2023-07-12 22:14:59 +08:00 committed by GitHub
parent 49a4be635d
commit cebe787bed
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 11 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
nerfstudio/engine/optimizers.py
View File

@ -159,3 +159,12 @@ class Optimizers:
"""
for k, v in loaded_state.items():
self.optimizers[k].load_state_dict(v)
def load_schedulers(self, loaded_state: Dict[str, Any]) -> None:
"""Helper to load the scheduler state from previous checkpoint
Args:
loaded_state: the state from the previous checkpoint
"""
for k, v in loaded_state.items():
self.schedulers[k].load_state_dict(v)

2
nerfstudio/engine/trainer.py
View File

@ -406,6 +406,7 @@ 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"])
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:
@ -439,6 +440,7 @@ class Trainer:
if hasattr(self.pipeline, "module")
else self.pipeline.state_dict(),
"optimizers": {k: v.state_dict() for (k, v) in self.optimizers.optimizers.items()},
"schedulers": {k: v.state_dict() for (k, v) in self.optimizers.schedulers.items()},
"scalers": self.grad_scaler.state_dict(),
},
ckpt_path,