diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index 21029e1b0..f3c18dab8 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/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) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index cf9f4bfec..44dd6d48e 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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), diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index a76354735..c0ca9501d 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/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 diff --git a/nerfstudio/engine/optimizers.py b/nerfstudio/engine/optimizers.py index 5bae109d0..e28de99bb 100644 --- a/nerfstudio/engine/optimizers.py +++ b/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) diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index bbfce5104..0951e6cf2 100644 --- a/nerfstudio/models/nerfacto.py +++ b/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)