diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index 03d048e3e..7c6627785 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/nerfstudio/cameras/camera_optimizers.py @@ -43,6 +43,11 @@ class CameraOptimizerConfig(InstantiateConfig): mode: Literal["off", "SO3xR3", "SE3"] = "off" """Pose optimization strategy to use. If enabled, we recommend SO3xR3.""" + trans_l2_penalty: float = 1e-3 + """L2 penalty on translation parameters.""" + + rot_l2_penalty: float = 1e-2 + """L2 penalty on rotation parameters.""" class CameraOptimizer(nn.Module): """Layer that modifies camera poses to be optimized as well as the field during training.""" @@ -105,3 +110,10 @@ class CameraOptimizer(nn.Module): correction_matrices = self(raybundle.camera_indices.squeeze()) raybundle.origins = raybundle.origins + correction_matrices[:, :3, 3] raybundle.directions = torch.bmm(correction_matrices[:, :3, :3], raybundle.directions[..., None]).squeeze() + + def get_loss_dict(self, loss_dict: dict) -> dict: + """ + Add a regularizer + """ + loss_dict['camera_opt_regularizer'] = self.pose_adjustment[:,:3].norm(dim=-1).mean() * self.config.trans_l2_penalty + loss_dict['camera_rot_regularizer'] = self.pose_adjustment[:,3:].norm(dim=-1).mean() * self.config.rot_l2_penalty \ No newline at end of file diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 96f7c4ccf..e25a78311 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -103,7 +103,7 @@ method_configs["nerfacto"] = TrainerConfig( "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000), }, "camera_opt": { - "optimizer": AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2), + "optimizer": AdamOptimizerConfig(lr=6e-4, eps=1e-8), "scheduler": ExponentialDecaySchedulerConfig(lr_final=6e-6, max_steps=200000), }, }, @@ -156,54 +156,53 @@ method_configs["nerfacto"] = TrainerConfig( # viewer=ViewerConfig(num_rays_per_chunk=1 << 15), # vis="viewer", # ) -# method_configs["nerfacto-huge"] = TrainerConfig( -# method_name="nerfacto", -# steps_per_eval_batch=500, -# steps_per_save=2000, -# max_num_iterations=100000, -# mixed_precision=True, -# pipeline=VanillaPipelineConfig( -# datamanager=VanillaDataManagerConfig( -# dataparser=NerfstudioDataParserConfig(), -# train_num_rays_per_batch=16384, -# eval_num_rays_per_batch=4096, -# camera_optimizer=CameraOptimizerConfig( -# mode="SO3xR3", -# optimizer=RAdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-3), -# scheduler=ExponentialDecaySchedulerConfig(lr_final=6e-5, max_steps=50000), -# ), -# ), -# model=NerfactoModelConfig( -# eval_num_rays_per_chunk=1 << 15, -# num_nerf_samples_per_ray=64, -# num_proposal_samples_per_ray=(512, 512), -# proposal_net_args_list=[ -# {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 512, "use_linear": False}, -# {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 7, "max_res": 2048, "use_linear": False}, -# ], -# hidden_dim=256, -# hidden_dim_color=256, -# appearance_embed_dim=32, -# features_per_level=4, -# base_res=32, -# max_res=8192, -# proposal_weights_anneal_max_num_iters=5000, -# log2_hashmap_size=21, -# ), -# ), -# optimizers={ -# "proposal_networks": { -# "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), -# "scheduler": None, -# }, -# "fields": { -# "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), -# "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=50000), -# }, -# }, -# viewer=ViewerConfig(num_rays_per_chunk=1 << 15), -# vis="viewer", -# ) +method_configs["nerfacto-huge"] = TrainerConfig( + method_name="nerfacto", + steps_per_eval_batch=500, + steps_per_save=2000, + max_num_iterations=100000, + mixed_precision=True, + pipeline=VanillaPipelineConfig( + datamanager=VanillaDataManagerConfig( + dataparser=NerfstudioDataParserConfig(), + train_num_rays_per_batch=16384, + eval_num_rays_per_batch=4096, + ), + model=NerfactoModelConfig( + eval_num_rays_per_chunk=1 << 15, + num_nerf_samples_per_ray=64, + num_proposal_samples_per_ray=(512, 512), + proposal_net_args_list=[ + {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 512, "use_linear": False}, + {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 7, "max_res": 2048, "use_linear": False}, + ], + hidden_dim=256, + hidden_dim_color=256, + appearance_embed_dim=32, + features_per_level=4, + base_res=32, + max_res=8192, + proposal_weights_anneal_max_num_iters=5000, + log2_hashmap_size=21, + ), + ), + optimizers={ + "proposal_networks": { + "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + }, + "fields": { + "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=50000), + }, + "camera_opt": { + "optimizer": RAdamOptimizerConfig(lr=6e-4, eps=1e-8), + "scheduler": ExponentialDecaySchedulerConfig(lr_final=6e-5, max_steps=50000), + } + }, + viewer=ViewerConfig(num_rays_per_chunk=1 << 15), + vis="viewer", +) # method_configs["depth-nerfacto"] = TrainerConfig( # method_name="depth-nerfacto", diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 2bbfc51e2..407e92fa6 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -127,7 +127,7 @@ class NerfactoModelConfig(ModelConfig): """Which implementation to use for the model.""" appearance_embed_dim: int = 32 """Dimension of the appearance embedding.""" - camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig() + camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig(mode="SO3xR3") """Config of the camera optimizer to use""" @@ -367,6 +367,8 @@ class NerfactoModel(Model): loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean( outputs["rendered_pred_normal_loss"] ) + # Add losses from the camera optimizer (like L2 regularization) + self.camera_optimizer.get_loss_dict(loss_dict) return loss_dict def get_image_metrics_and_images(