Add average_init_density to improve robustness of nerfacto training

Author: Jianbo Ye

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
 

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
 

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 = 0.01
+    """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)