From e4a0050d0d61022530204e1644998afd0f77fb2d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E3=83=95=E3=83=A9=E3=83=8F=E3=83=86=E3=82=A3=E3=80=80?= =?UTF-8?q?=E4=BB=81?= Date: Fri, 27 Oct 2023 07:15:58 +0900 Subject: [PATCH] Add gradient scaling option to more methods (#2555) add gradient scaling to instant-ngp, mipnerf, tensorf, vanilla-nerf --- nerfstudio/models/instant_ngp.py | 6 +++++- nerfstudio/models/mipnerf.py | 6 +++++- nerfstudio/models/tensorf.py | 6 +++++- nerfstudio/models/vanilla_nerf.py | 8 +++++++- 4 files changed, 22 insertions(+), 4 deletions(-) diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py index c67d45db5..48cf1a95a 100644 --- a/nerfstudio/models/instant_ngp.py +++ b/nerfstudio/models/instant_ngp.py @@ -37,7 +37,7 @@ from nerfstudio.engine.callbacks import ( from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.field_components.spatial_distortions import SceneContraction from nerfstudio.fields.nerfacto_field import NerfactoField -from nerfstudio.model_components.losses import MSELoss +from nerfstudio.model_components.losses import MSELoss, scale_gradients_by_distance_squared from nerfstudio.model_components.ray_samplers import VolumetricSampler from nerfstudio.model_components.renderers import ( AccumulationRenderer, @@ -78,6 +78,8 @@ class InstantNGPModelConfig(ModelConfig): """How far along ray to start sampling.""" far_plane: float = 1e3 """How far along ray to stop sampling.""" + use_gradient_scaling: bool = False + """Use gradient scaler where the gradients are lower for points closer to the camera.""" use_appearance_embedding: bool = False """Whether to use an appearance embedding.""" background_color: Literal["random", "black", "white"] = "random" @@ -187,6 +189,8 @@ class NGPModel(Model): ) field_outputs = self.field(ray_samples) + if self.config.use_gradient_scaling: + field_outputs = scale_gradients_by_distance_squared(field_outputs, ray_samples) # accumulation packed_info = nerfacc.pack_info(ray_indices, num_rays) diff --git a/nerfstudio/models/mipnerf.py b/nerfstudio/models/mipnerf.py index 67ca63251..cf48b08a7 100644 --- a/nerfstudio/models/mipnerf.py +++ b/nerfstudio/models/mipnerf.py @@ -29,7 +29,7 @@ from nerfstudio.cameras.rays import RayBundle from nerfstudio.field_components.encodings import NeRFEncoding from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.fields.vanilla_nerf_field import NeRFField -from nerfstudio.model_components.losses import MSELoss +from nerfstudio.model_components.losses import MSELoss, scale_gradients_by_distance_squared from nerfstudio.model_components.ray_samplers import PDFSampler, UniformSampler from nerfstudio.model_components.renderers import ( AccumulationRenderer, @@ -109,6 +109,8 @@ class MipNerfModel(Model): # First pass: field_outputs_coarse = self.field.forward(ray_samples_uniform) + if self.config.use_gradient_scaling: + field_outputs_coarse = scale_gradients_by_distance_squared(field_outputs_coarse, ray_samples_uniform) weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY]) rgb_coarse = self.renderer_rgb( rgb=field_outputs_coarse[FieldHeadNames.RGB], @@ -122,6 +124,8 @@ class MipNerfModel(Model): # Second pass: field_outputs_fine = self.field.forward(ray_samples_pdf) + if self.config.use_gradient_scaling: + field_outputs_fine = scale_gradients_by_distance_squared(field_outputs_fine, ray_samples_pdf) weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY]) rgb_fine = self.renderer_rgb( rgb=field_outputs_fine[FieldHeadNames.RGB], diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py index 9cc9c7ccf..53d986337 100644 --- a/nerfstudio/models/tensorf.py +++ b/nerfstudio/models/tensorf.py @@ -43,7 +43,7 @@ from nerfstudio.field_components.encodings import ( ) from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.fields.tensorf_field import TensoRFField -from nerfstudio.model_components.losses import MSELoss, tv_loss +from nerfstudio.model_components.losses import MSELoss, tv_loss, scale_gradients_by_distance_squared from nerfstudio.model_components.ray_samplers import PDFSampler, UniformSampler from nerfstudio.model_components.renderers import ( AccumulationRenderer, @@ -92,6 +92,8 @@ class TensoRFModelConfig(ModelConfig): """Regularization method used in tensorf paper""" camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig(mode="SO3xR3") """Config of the camera optimizer to use""" + use_gradient_scaling: bool = False + """Use gradient scaler where the gradients are lower for points closer to the camera.""" background_color: Literal["random", "last_sample", "black", "white"] = "white" """Whether to randomize the background color.""" @@ -296,6 +298,8 @@ class TensoRFModel(Model): field_outputs_fine = self.field.forward( ray_samples_pdf, mask=acc_mask, bg_color=colors.WHITE.to(weights.device) ) + if self.config.use_gradient_scaling: + field_outputs_fine = scale_gradients_by_distance_squared(field_outputs_fine, ray_samples_pdf) weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY]) diff --git a/nerfstudio/models/vanilla_nerf.py b/nerfstudio/models/vanilla_nerf.py index e95c8dfa4..4f4dbd8f6 100644 --- a/nerfstudio/models/vanilla_nerf.py +++ b/nerfstudio/models/vanilla_nerf.py @@ -33,7 +33,7 @@ from nerfstudio.field_components.encodings import NeRFEncoding from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind from nerfstudio.fields.vanilla_nerf_field import NeRFField -from nerfstudio.model_components.losses import MSELoss +from nerfstudio.model_components.losses import MSELoss, scale_gradients_by_distance_squared from nerfstudio.model_components.ray_samplers import PDFSampler, UniformSampler from nerfstudio.model_components.renderers import ( AccumulationRenderer, @@ -58,6 +58,8 @@ class VanillaModelConfig(ModelConfig): """Specifies whether or not to include ray warping based on time.""" temporal_distortion_params: Dict[str, Any] = to_immutable_dict({"kind": TemporalDistortionKind.DNERF}) """Parameters to instantiate temporal distortion with""" + use_gradient_scaling: bool = False + """Use gradient scaler where the gradients are lower for points closer to the camera.""" background_color: Literal["random", "last_sample", "black", "white"] = "white" """Whether to randomize the background color.""" @@ -154,6 +156,8 @@ class NeRFModel(Model): # coarse field: field_outputs_coarse = self.field_coarse.forward(ray_samples_uniform) + if self.config.use_gradient_scaling: + field_outputs_coarse = scale_gradients_by_distance_squared(field_outputs_coarse, ray_samples_uniform) weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY]) rgb_coarse = self.renderer_rgb( rgb=field_outputs_coarse[FieldHeadNames.RGB], @@ -172,6 +176,8 @@ class NeRFModel(Model): # fine field: field_outputs_fine = self.field_fine.forward(ray_samples_pdf) + if self.config.use_gradient_scaling: + field_outputs_fine = scale_gradients_by_distance_squared(field_outputs_fine, ray_samples_pdf) weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY]) rgb_fine = self.renderer_rgb( rgb=field_outputs_fine[FieldHeadNames.RGB],