From 57d0e149e9a4b7f548f7058be8509c2f5f98e23d Mon Sep 17 00:00:00 2001 From: jb-ye Date: Sat, 27 Jan 2024 05:47:47 +0000 Subject: [PATCH] Add average_init_density to improve robustness of nerfacto training --- nerfstudio/configs/method_configs.py | 3 +++ nerfstudio/fields/density_fields.py | 4 +++- nerfstudio/fields/nerfacto_field.py | 4 +++- nerfstudio/models/nerfacto.py | 5 +++++ 4 files changed, 14 insertions(+), 2 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 9a7d31c10c6..b3faf2d9eec 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -97,6 +97,7 @@ ), model=NerfactoModelConfig( eval_num_rays_per_chunk=1 << 15, + average_init_density=0.01, camera_optimizer=CameraOptimizerConfig(mode="SO3xR3"), ), ), @@ -140,6 +141,7 @@ max_res=4096, proposal_weights_anneal_max_num_iters=5000, log2_hashmap_size=21, + average_init_density=0.01, camera_optimizer=CameraOptimizerConfig(mode="SO3xR3"), ), ), @@ -187,6 +189,7 @@ max_res=8192, proposal_weights_anneal_max_num_iters=5000, log2_hashmap_size=21, + average_init_density=0.01, camera_optimizer=CameraOptimizerConfig(mode="SO3xR3"), ), ), diff --git a/nerfstudio/fields/density_fields.py b/nerfstudio/fields/density_fields.py index 0837e63820f..42fdc13dc55 100644 --- a/nerfstudio/fields/density_fields.py +++ b/nerfstudio/fields/density_fields.py @@ -56,6 +56,7 @@ def __init__( base_res: int = 16, log2_hashmap_size: int = 18, features_per_level: int = 2, + average_init_density: float = 1.0, implementation: Literal["tcnn", "torch"] = "tcnn", ) -> None: super().__init__() @@ -66,6 +67,7 @@ def __init__( self.register_buffer("max_res", torch.tensor(max_res)) self.register_buffer("num_levels", torch.tensor(num_levels)) self.register_buffer("log2_hashmap_size", torch.tensor(log2_hashmap_size)) + self.register_buffer("average_init_density", torch.tensor(average_init_density)) self.encoding = HashEncoding( num_levels=num_levels, @@ -111,7 +113,7 @@ def get_density(self, ray_samples: RaySamples) -> Tuple[Tensor, None]: # Rectifying the density with an exponential is much more stable than a ReLU or # softplus, because it enables high post-activation (float32) density outputs # from smaller internal (float16) parameters. - density = trunc_exp(density_before_activation) + density = self.average_init_density * trunc_exp(density_before_activation) density = density * selector[..., None] return density, None diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index 910bba8465c..1cef23ab784 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/nerfstudio/fields/nerfacto_field.py @@ -95,6 +95,7 @@ def __init__( use_pred_normals: bool = False, use_average_appearance_embedding: bool = False, spatial_distortion: Optional[SpatialDistortion] = None, + average_init_density: float = 1.0, implementation: Literal["tcnn", "torch"] = "tcnn", ) -> None: super().__init__() @@ -105,6 +106,7 @@ def __init__( self.register_buffer("max_res", torch.tensor(max_res)) self.register_buffer("num_levels", torch.tensor(num_levels)) self.register_buffer("log2_hashmap_size", torch.tensor(log2_hashmap_size)) + self.register_buffer("average_init_density", torch.tensor(average_init_density)) self.spatial_distortion = spatial_distortion self.num_images = num_images @@ -218,7 +220,7 @@ def get_density(self, ray_samples: RaySamples) -> Tuple[Tensor, Tensor]: # Rectifying the density with an exponential is much more stable than a ReLU or # softplus, because it enables high post-activation (float32) density outputs # from smaller internal (float16) parameters. - density = trunc_exp(density_before_activation.to(positions)) + density = self.average_init_density * trunc_exp(density_before_activation.to(positions)) density = density * selector[..., None] return density, base_mlp_out diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index b1898f3d66b..47fc42e31f9 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -124,6 +124,8 @@ class NerfactoModelConfig(ModelConfig): """Which implementation to use for the model.""" appearance_embed_dim: int = 32 """Dimension of the appearance embedding.""" + average_init_density: float = 1.0 + """Average initial density output from MLP. """ camera_optimizer: CameraOptimizerConfig = field(default_factory=lambda: CameraOptimizerConfig(mode="SO3xR3")) """Config of the camera optimizer to use""" @@ -162,6 +164,7 @@ def populate_modules(self): use_pred_normals=self.config.predict_normals, use_average_appearance_embedding=self.config.use_average_appearance_embedding, appearance_embedding_dim=self.config.appearance_embed_dim, + average_init_density=self.config.average_init_density, implementation=self.config.implementation, ) @@ -179,6 +182,7 @@ def populate_modules(self): self.scene_box.aabb, spatial_distortion=scene_contraction, **prop_net_args, + average_init_density=self.config.average_init_density, implementation=self.config.implementation, ) self.proposal_networks.append(network) @@ -190,6 +194,7 @@ def populate_modules(self): self.scene_box.aabb, spatial_distortion=scene_contraction, **prop_net_args, + average_init_density=self.config.average_init_density, implementation=self.config.implementation, ) self.proposal_networks.append(network)