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,34 @@ 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.pose_opt and "pose_adjust" in ckpt:
+                pose_adjust = self.pose_adjust.module if world_size > 1 else self.pose_adjust
+                pose_adjust.load_state_dict(ckpt["pose_adjust"])
+            if cfg.app_opt and "app_module" in ckpt:
+                app_module = self.app_module.module if world_size > 1 else self.app_module
+                app_module.load_state_dict(ckpt["app_module"])
+            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)
+                if cfg.pose_opt and "pose_optimizer_state_dicts" in ckpt:
+                    for optimizer, opt_state_dict in zip(self.pose_optimizers, ckpt["pose_optimizer_state_dicts"]):
+                        optimizer.load_state_dict(opt_state_dict)
+                if cfg.app_opt and "app_optimizer_state_dicts" in ckpt:
+                    for optimizer, opt_state_dict in zip(self.app_optimizers, ckpt["app_optimizer_state_dicts"]):
+                        optimizer.load_state_dict(opt_state_dict)
+                if cfg.use_bilateral_grid and "bil_grid_optimizer_state_dicts" in ckpt:
+                    for optimizer, opt_state_dict in zip(self.bil_grid_optimizers, ckpt["bil_grid_optimizer_state_dicts"]):
+                        optimizer.load_state_dict(opt_state_dict)
+                step = ckpt["step"] - 1
+
         if not cfg.use_progress_bar:
             logging.info("Training started")
 
@@ -746,13 +797,14 @@ def train(self):
                         packed=cfg.packed,
                     )
                 elif isinstance(self.cfg.strategy, MCMCStrategy):
+                    means_lr = schedulers[0].get_last_lr()[0] if cfg.use_scheduler else self.optimizers["means"].param_groups[0]["lr"]
                     self.cfg.strategy.step_post_backward(
                         params=self.splats,
                         optimizers=self.optimizers,
                         state=self.strategy_state,
                         step=step,
                         info=info,
-                        lr=schedulers[0].get_last_lr()[0],
+                        lr=means_lr,
                     )
                 else:
                     assert_never(self.cfg.strategy)
@@ -763,7 +815,16 @@ 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:
+                    data["pose_optimizer_state_dicts"] = [optimizer.state_dict() for optimizer in self.pose_optimizers]
+                if cfg.app_opt:
+                    data["app_optimizer_state_dicts"] = [optimizer.state_dict() for optimizer in self.app_optimizers]
+                if cfg.use_bilateral_grid:
+                    data["bil_grid_optimizer_state_dicts"] = [optimizer.state_dict() for optimizer in self.bil_grid_optimizers]
 
                 if cfg.pose_opt:
                     if world_size > 1:
@@ -784,13 +845,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/gaussianSplattingPhotogrammetry.mg

@@ -1,20 +1,20 @@
 {
     "header": {
         "releaseVersion": "2026.1.0+develop",
-        "fileVersion": "2.0",
+        "fileVersion": "2.1",
         "nodesVersions": {
-            "CameraInit": "12.0",
+            "CameraInit": "12.1",
             "ExportImages": "1.1",
             "FeatureExtraction": "1.3",
-            "FeatureMatching": "2.0",
-            "GaussianSplattingOptim": "1.0",
+            "FeatureMatching": "2.1",
+            "GaussianSplattingOptim": "1.1",
             "GaussianSplattingRender": "1.0",
             "ImageMatching": "2.0",
             "IntrinsicsTransforming": "1.1",
             "RelativePoseEstimating": "3.1",
-            "SfMBootStrapping": "4.1",
+            "SfMBootStrapping": "4.2",
             "SfMColorizing": "1.0",
-            "SfMExpanding": "2.3",
+            "SfMExpanding": "2.4",
             "SfMTransform": "3.2",
             "TracksBuilding": "1.0"
         },
@@ -118,9 +118,14 @@
             ],
             "inputs": {
                 "sfm_file": "{ExportImages_2.outputSfMData}",
-                "image_alpha": true,
+                "image_alpha": "{GaussianSplattingOptim_1.image_alpha}",
                 "resume_ckpt": "{GaussianSplattingOptim_1.model}",
-                "max_epochs": 100
+                "strategy": "{GaussianSplattingOptim_1.strategy}",
+                "max_epochs": 100,
+                "refine_start_iter": "{GaussianSplattingOptim_1.refine_start_iter}",
+                "refine_stop_iter": "{GaussianSplattingOptim_1.refine_stop_iter}",
+                "reset_every": "{GaussianSplattingOptim_1.reset_every}",
+                "refine_every": "{GaussianSplattingOptim_1.refine_every}"
             }
         },
         "GaussianSplattingOptim_3": {
@@ -131,9 +136,14 @@
             ],
             "inputs": {
                 "sfm_file": "{ExportImages_3.outputSfMData}",
-                "image_alpha": true,
+                "image_alpha": "{GaussianSplattingOptim_2.image_alpha}",
                 "resume_ckpt": "{GaussianSplattingOptim_2.model}",
-                "max_epochs": 100
+                "strategy": "{GaussianSplattingOptim_2.strategy}",
+                "max_epochs": 100,
+                "refine_start_iter": "{GaussianSplattingOptim_2.refine_start_iter}",
+                "refine_stop_iter": "{GaussianSplattingOptim_2.refine_stop_iter}",
+                "reset_every": "{GaussianSplattingOptim_2.reset_every}",
+                "refine_every": "{GaussianSplattingOptim_2.refine_every}"
             }
         },
         "GaussianSplattingRender_1": {
@@ -265,4 +275,4 @@
             }
         }
     }
-}
+}

meshroom/mrGSplat/GaussianSplattingOptim.py

@@ -1,4 +1,4 @@
-__version__ = "1.0"
+__version__ = "1.1"
 
 from meshroom.core import desc
 
@@ -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',