From cebe787bedffd4966fdcb067682da64e168b5565 Mon Sep 17 00:00:00 2001
From: hyperzy <50564821+hyperzy@users.noreply.github.com>
Date: Wed, 12 Jul 2023 22:14:59 +0800
Subject: [PATCH] Add saving/loading scheduler state_dict (#2203)

- Add saving/loading scheduler stat_dict
---
 nerfstudio/engine/optimizers.py | 9 +++++++++
 nerfstudio/engine/trainer.py    | 2 ++
 2 files changed, 11 insertions(+)

diff --git a/nerfstudio/engine/optimizers.py b/nerfstudio/engine/optimizers.py
index 3ff693685..3e147dcdc 100644
--- a/nerfstudio/engine/optimizers.py
+++ b/nerfstudio/engine/optimizers.py
@@ -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)
diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py
index 0e194b631..a89e4801c 100644
--- a/nerfstudio/engine/trainer.py
+++ b/nerfstudio/engine/trainer.py
@@ -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,