tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

Add average_init_density to improve robustness of nerfacto training (#2834)

Browse Source
This commit is contained in:
J.Y 2024-01-28 17:09:27 -05:00 committed by GitHub
parent fb1fee1d8f
commit d8b517d6a5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 14 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
nerfstudio/configs/method_configs.py
View File

@ -97,6 +97,7 @@ method_configs["nerfacto"] = TrainerConfig(
),
model=NerfactoModelConfig(
eval_num_rays_per_chunk=1 << 15,
average_init_density=0.01,
camera_optimizer=CameraOptimizerConfig(mode="SO3xR3"),
),
),
@ -140,6 +141,7 @@ method_configs["nerfacto-big"] = TrainerConfig(
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 @@ method_configs["nerfacto-huge"] = TrainerConfig(
max_res=8192,
proposal_weights_anneal_max_num_iters=5000,
log2_hashmap_size=21,
average_init_density=0.01,
camera_optimizer=CameraOptimizerConfig(mode="SO3xR3"),
),
),

4
nerfstudio/fields/density_fields.py
View File

@ -56,12 +56,14 @@ class HashMLPDensityField(Field):
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__()
self.register_buffer("aabb", aabb)
self.spatial_distortion = spatial_distortion
self.use_linear = use_linear
self.average_init_density = average_init_density
self.register_buffer("max_res", torch.tensor(max_res))
self.register_buffer("num_levels", torch.tensor(num_levels))
@ -111,7 +113,7 @@ class HashMLPDensityField(Field):
# 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

4
nerfstudio/fields/nerfacto_field.py
View File

@ -95,6 +95,7 @@ class NerfactoField(Field):
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__()
@ -116,6 +117,7 @@ class NerfactoField(Field):
self.use_pred_normals = use_pred_normals
self.pass_semantic_gradients = pass_semantic_gradients
self.base_res = base_res
self.average_init_density = average_init_density
self.step = 0
self.direction_encoding = SHEncoding(
@ -218,7 +220,7 @@ class NerfactoField(Field):
# 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

5
nerfstudio/models/nerfacto.py
View File

@ -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 @@ class NerfactoModel(Model):
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 @@ class NerfactoModel(Model):
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 @@ class NerfactoModel(Model):
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)