Correct the model retrieval

Author: CedricThebault

gsplatInterface/trainer.py

@@ -39,6 +39,10 @@
 from typing import Literal, assert_never
 from utils import AppearanceOptModule, CameraOptModule, knn, rgb_to_sh, set_random_seed
 
+from numpy.core.multiarray import scalar as npscalar
+from numpy.dtypes import Float64DType as npFloat64DType
+from numpy import dtype as npdtype
+
 from gsplat import export_splats
 from gsplat.distributed import cli
 from gsplat.optimizers import SelectiveAdam
@@ -58,6 +62,9 @@ class Config:
     # Path to the .pt files to load and resume training
     resume_ckpt: Optional[str] = ""
 
+    # Whether to retrieve optimizer state from the checkpoint
+    retrieve_optimizer_state: bool = True
+
     # Path to the dataset
     sfm_file: str = "sfm.json"
     image_alpha: bool = False
@@ -127,6 +134,9 @@ class Config:
     # Use random background for training to discourage transparency
     random_bkgd: bool = False
 
+    # Whether to use schedulers for the learning rates (only for MCMC)
+    use_scheduler: bool = True
+
     # LR for 3D point positions
     means_lr: float = 1.6e-4
     # LR for Gaussian scale factors
@@ -222,6 +232,7 @@ def create_splats_with_optimizers(
     device: str = "cuda",
     world_rank: int = 0,
     world_size: int = 1,
+    strategy: Union[DefaultStrategy, MCMCStrategy] = DefaultStrategy,
 ) -> Tuple[torch.nn.ParameterDict, Dict[str, torch.optim.Optimizer]]:
 
     #initialize from sfm
@@ -328,6 +339,12 @@ def __init__(
 
         logging.info(f"Scene scale: {self.scene_scale}")
 
+        retrieve_optimizer_state = cfg.retrieve_optimizer_state and cfg.resume_ckpt != ""
+        if retrieve_optimizer_state and isinstance(cfg.strategy, MCMCStrategy):
+            # In case of MCMC, the scheduler reduces the effective learning rate by 100
+            # This line updates the learning rate accordingly
+            cfg.means_lr = 0.01 * cfg.means_lr
+
         # Model
         feature_dim = 32 if cfg.app_opt else None
         self.splats, self.optimizers = create_splats_with_optimizers(
@@ -349,6 +366,7 @@ def __init__(
             device=self.device,
             world_rank=world_rank,
             world_size=world_size,
+            strategy=cfg.strategy,
         )
         logging.info(f"Model initialized. Number of GS: {len(self.splats['means'])}")
 
@@ -517,37 +535,42 @@ def train(self):
         #Compute effective steps
         max_steps = cfg.max_epochs * len(self.trainset)
 
-        schedulers = [
-            # means has a learning rate schedule, that end at 0.01 of the initial value
-            torch.optim.lr_scheduler.ExponentialLR(
-                self.optimizers["means"], gamma=0.01 ** (1.0 / max_steps)
-            ),
-        ]
+        schedulers = []
+        if isinstance(cfg.strategy, MCMCStrategy) and cfg.use_scheduler:
 
-        if cfg.pose_opt:
-            # pose optimization has a learning rate schedule
-            schedulers.append(
+            logging.info("Create schedulers")
+
+            schedulers = [
+                # means has a learning rate schedule that ends at 0.01 of the initial value
                 torch.optim.lr_scheduler.ExponentialLR(
-                    self.pose_optimizers[0], gamma=0.01 ** (1.0 / max_steps)
+                    self.optimizers["means"], gamma=0.01 ** (1.0 / max_steps)
+                ),
+            ]
+
+            if cfg.pose_opt:
+                # pose optimization has a learning rate schedule
+                schedulers.append(
+                    torch.optim.lr_scheduler.ExponentialLR(
+                        self.pose_optimizers[0], gamma=0.01 ** (1.0 / max_steps)
+                    )
                 )
-            )
 
-        if cfg.use_bilateral_grid:
-            # bilateral grid has a learning rate schedule. Linear warmup for 1000 steps.
-            schedulers.append(
-                torch.optim.lr_scheduler.ChainedScheduler(
-                    [
-                        torch.optim.lr_scheduler.LinearLR(
-                            self.bil_grid_optimizers[0],
-                            start_factor=0.01,
-                            total_iters=1000,
-                        ),
-                        torch.optim.lr_scheduler.ExponentialLR(
-                            self.bil_grid_optimizers[0], gamma=0.01 ** (1.0 / max_steps)
-                        ),
-                    ]
+            if cfg.use_bilateral_grid:
+                # bilateral grid has a learning rate schedule. Linear warmup for 1000 steps.
+                schedulers.append(
+                    torch.optim.lr_scheduler.ChainedScheduler(
+                        [
+                            torch.optim.lr_scheduler.LinearLR(
+                                self.bil_grid_optimizers[0],
+                                start_factor=0.01,
+                                total_iters=1000,
+                            ),
+                            torch.optim.lr_scheduler.ExponentialLR(
+                                self.bil_grid_optimizers[0], gamma=0.01 ** (1.0 / max_steps)
+                            ),
+                        ]
+                    )
                 )
-            )
 
         trainloader = torch.utils.data.DataLoader(
             self.trainset,
@@ -560,6 +583,19 @@ def train(self):
 
         step = -1
 
+        if cfg.resume_ckpt != "":
+            torch.serialization.add_safe_globals([npscalar, npdtype, npFloat64DType])
+            logging.info("Loading Gaussians from checkpoint")
+            ckpt = torch.load(cfg.resume_ckpt, map_location=device, weights_only=True)
+            for k in self.splats.keys():
+                self.splats[k].data = torch.cat([ckpt["splats"][k]])
+            if  cfg.retrieve_optimizer_state:
+                logging.info("Loading optimizer state from checkpoint")
+                self.strategy_state = ckpt["strategy_state"]
+                for name, opt_state_dict in ckpt["optimizer_state_dicts"].items():
+                    self.optimizers[name].load_state_dict(opt_state_dict)
+                step = ckpt["step"] - 1
+
         if not cfg.use_progress_bar:
             logging.info("Training started")
 
@@ -763,7 +799,10 @@ def train(self):
             if epoch in [i - 1 for i in cfg.save_epochs] or epoch == cfg.max_epochs - 1:
 
                 path = f"{self.ckpt_dir}/ckpt_{epoch}_rank{self.world_rank}.pt"
-                data = {"step": step, "splats": self.splats.state_dict()}
+                data = {"step": step,
+                        "splats": self.splats.state_dict(),
+                        "strategy_state": self.strategy_state,
+                        "optimizer_state_dicts": {name: opt.state_dict() for name, opt in self.optimizers.items()}}
 
                 if cfg.pose_opt:
                     if world_size > 1:
@@ -784,13 +823,6 @@ def main(local_rank: int, world_rank, world_size: int, cfg: Config):
     #Create object for computation
     runner = Runner(local_rank, world_rank, world_size, cfg)
 
-    #if a checkpoint was passed, will load the gaussians from there
-    if cfg.resume_ckpt != "":
-        logging.info("Loading from checkpoint")
-        ckpt = torch.load(cfg.resume_ckpt, map_location=runner.device, weights_only=True)
-        for k in runner.splats.keys():
-            runner.splats[k].data = torch.cat([ckpt["splats"][k]])
-
     #Launch training
     runner.train()
 

gsplatInterface/viewer.py

@@ -18,6 +18,9 @@
 
 from datasets.sfm.sceneManager import PoseParser
 # from cameraPosesParser import CameraParser
+from numpy.core.multiarray import scalar as npscalar
+from numpy.dtypes import Float64DType as npFloat64DType
+from numpy import dtype as npdtype
 
 
 def get_cameras_from_sfm(sfmFile):
@@ -56,7 +59,8 @@ def get_cameras_from_sfm(sfmFile):
 def main(local_rank: int, world_rank, world_size: int, args):
     torch.manual_seed(42)
     device = torch.device("cuda", local_rank)
-    
+    torch.serialization.add_safe_globals([npscalar, npdtype, npFloat64DType])
+
     means, quats, scales, opacities, sh0, shN = [], [], [], [], [], []
     for ckpt_path in args.ckpt:
         ckpt = torch.load(ckpt_path, map_location=device, weights_only=True)["splats"]

meshroom/mrGSplat/GaussianSplattingOptim.py

@@ -73,12 +73,14 @@ def buildCommandLine(self, chunk):
         if node.revisedOpacity.value:
             cmdLine += " --revised_opacity"
 
+        if not node.retrieveOptimState.value:
+            cmdLine += " --no-retrieve_optimizer_state"
+
         if node.useProgressBar.value:
             cmdLine += f" --use_progress_bar"
 
         node.nodeDesc.commandLine = cmdLine
 
-
         return super().buildCommandLine(chunk)
 
     inputs = [
@@ -101,6 +103,14 @@ def buildCommandLine(self, chunk):
             description="Resume from Model",
             value="",
         ),
+        desc.BoolParam(
+            name="retrieveOptimState",
+            label="Retrieve Optimizer State",
+            description="Whether to retrieve optimizer state from the checkpoint.",
+            value=True,
+            commandLineGroup=None,
+            enabled=lambda node: node.resume_ckpt.value != ""
+        ),
         desc.ChoiceParam(
             name='strategy',
             label='Densification Strategy',