fixed direction sampling, separate optimizers, playing with LR

nerfstudio/cameras/camera_optimizers.py

@@ -53,6 +53,14 @@ class CameraOptimizerConfig(cfg.InstantiateConfig):
     optimizer: AdamOptimizerConfig = AdamOptimizerConfig(lr=6e-4, eps=1e-15)
     """ADAM parameters for camera optimization."""
 
+    blur_opt: AdamOptimizerConfig = AdamOptimizerConfig(lr=6e-4, eps=1e-15)
+    """Optimizer for blur params"""
+
+    blur_sched: SchedulerConfig = SchedulerConfig(max_steps=10000, lr_final=1e-3)
+    """scheduler for blur optimization"""
+
+    blur_group: tyro.conf.Suppress[str] = "blur_opt"
+
     scheduler: SchedulerConfig = SchedulerConfig(max_steps=10000, lr_final=1e-5)
     """Learning rate scheduler for camera optimizer.."""
 
@@ -85,6 +93,10 @@ class CameraOptimizer(nn.Module):
         elif self.config.mode in ("SO3xR3", "SE3"):
             self.pose_adjustment = torch.nn.Parameter(torch.zeros((num_cameras, 6), device=device))
             self.blur_adjustment = torch.nn.Parameter(torch.zeros((num_cameras, 6), device=device))
+            n = 5
+            b = 1500
+            ids = [torch.ones((b, 1), device=device, dtype=torch.float32) * delta for delta in torch.linspace(-1, 1, n)]
+            self.times = torch.concat(ids, dim=0).to(device)
         else:
             assert_never(self.config.mode)
 
@@ -136,11 +148,11 @@ class CameraOptimizer(nn.Module):
         Samples the blur parameters and returns transformation matrices for the origins and directions
         """
         self.i += 1
-        # blur_adjs = self.blur_adjustment[indices, :] * torch.rand((indices.shape[0], 1), device=self.device)
-        blur_adjs = self.blur_adjustment[indices, :] * torch.normal(
-            torch.zeros((indices.shape[0], 1), device=self.device),  # mean
-            torch.ones((indices.shape[0], 1), device=self.device),  # std
-        )
+        blur_adjs = self.blur_adjustment[indices, :] * self.times
+        # blur_adjs = self.blur_adjustment[indices, :] * torch.normal(
+        #     torch.zeros((indices.shape[0], 1), device=self.device),  # mean
+        #     torch.ones((indices.shape[0], 1), device=self.device),  # std
+        # )
         # if self.i % 500 == 0:
         #     # print the directions
         #     print(self.blur_adjustment)

nerfstudio/configs/method_configs.py

@@ -68,8 +68,10 @@ method_configs["nerfacto"] = Config(
             eval_num_rays_per_batch=4096,
             camera_optimizer=CameraOptimizerConfig(
                 mode="SO3xR3",
-                optimizer=AdamOptimizerConfig(lr=1e-2, eps=1e-8, weight_decay=0),
-                scheduler=SchedulerConfig(max_steps=10000, lr_final=1e-3),
+                optimizer=AdamOptimizerConfig(lr=1e-3, eps=1e-8, weight_decay=1e-3),
+                scheduler=SchedulerConfig(max_steps=3000, lr_final=1e-7),
+                blur_opt=AdamOptimizerConfig(lr=1e-5, weight_decay=0),
+                blur_sched=SchedulerConfig(max_steps=20000, lr_final=1e-6),
             ),
         ),
         model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15),

nerfstudio/data/datamanagers/base_datamanager.py

@@ -436,7 +436,8 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         camera_opt_params = list(self.train_camera_optimizer.parameters())
         if self.config.camera_optimizer.mode != "off":
             assert len(camera_opt_params) > 0
-            param_groups[self.config.camera_optimizer.param_group] = camera_opt_params
+            param_groups[self.config.camera_optimizer.param_group] = [self.train_camera_optimizer.pose_adjustment]
+            param_groups[self.config.camera_optimizer.blur_group] = [self.train_camera_optimizer.blur_adjustment]
         else:
             assert len(camera_opt_params) == 0
 

nerfstudio/engine/optimizers.py

@@ -81,6 +81,10 @@ def setup_optimizers(config: base_config.Config, param_groups: Dict[str, List[Pa
             "optimizer": config.pipeline.datamanager.camera_optimizer.optimizer,
             "scheduler": config.pipeline.datamanager.camera_optimizer.scheduler,
         }
+        optimizer_config[camera_optimizer_config.blur_group] = {
+            "optimizer": config.pipeline.datamanager.camera_optimizer.blur_opt,
+            "scheduler": config.pipeline.datamanager.camera_optimizer.blur_sched,
+        }
     return Optimizers(optimizer_config, param_groups)
 
 

nerfstudio/models/nerfacto.py

@@ -217,8 +217,9 @@ class NerfactoModel(Model):
 
         rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
         s = 1500
-        if rgb.shape[0] == s * 5:
-            rgb = (1 / 5.0) * (
+        n = 5.0
+        if rgb.shape[0] == s * int(n):
+            rgb = (1 / n) * (
                 rgb[:s, :] + rgb[s : 2 * s, :] + rgb[2 * s : 3 * s, :] + rgb[3 * s : 4 * s, :] + rgb[4 * s : 5 * s, :]
             )
         depth = self.renderer_depth(weights=weights, ray_samples=ray_samples)