diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index a92c042bf..113ecadbd 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -41,7 +41,7 @@ from nerfstudio.data.dataparsers.sitcoms3d_dataparser import Sitcoms3DDataParser from nerfstudio.data.datasets.depth_dataset import DepthDataset from nerfstudio.data.datasets.sdf_dataset import SDFDataset from nerfstudio.data.datasets.semantic_dataset import SemanticDataset -from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig +from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig, AdamWOptimizerConfig from nerfstudio.engine.schedulers import ( CosineDecaySchedulerConfig, ExponentialDecaySchedulerConfig, @@ -57,6 +57,7 @@ from nerfstudio.models.mipnerf import MipNerfModel from nerfstudio.models.nerfacto import NerfactoModelConfig from nerfstudio.models.neus import NeuSModelConfig from nerfstudio.models.neus_facto import NeuSFactoModelConfig +from nerfstudio.models.neuralangelo import NeuralangeloModelConfig from nerfstudio.models.semantic_nerfw import SemanticNerfWModelConfig from nerfstudio.models.tensorf import TensoRFModelConfig from nerfstudio.models.vanilla_nerf import NeRFModel, VanillaModelConfig @@ -79,6 +80,7 @@ descriptions = { "generfacto": "Generative Text to NeRF model", "neus": "Implementation of NeuS. (slow)", "neus-facto": "Implementation of NeuS-Facto. (slow)", + "neuralangelo": "Implementation of Neuralangelo. (slow)", } method_configs["nerfacto"] = TrainerConfig( @@ -575,6 +577,65 @@ method_configs["neus-facto"] = TrainerConfig( vis="viewer", ) +method_configs["neuralangelo"] = TrainerConfig( + method_name="neuralangelo", + steps_per_eval_image=5000, + steps_per_eval_batch=5000, + steps_per_save=20000, + steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images + max_num_iterations=500_001, + mixed_precision=False, + pipeline=VanillaPipelineConfig( + datamanager=VanillaDataManagerConfig( + _target=VanillaDataManager[SDFDataset], + dataparser=SDFStudioDataParserConfig(), + train_num_rays_per_batch=512, + eval_num_rays_per_batch=512, + camera_optimizer=CameraOptimizerConfig( + mode="off", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) + ), + ), + model=NeuralangeloModelConfig( + sdf_field=SDFFieldConfig( + use_grid_feature=True, + num_layers=1, + num_layers_color=4, + hidden_dim=256, + hidden_dim_color=256, + geometric_init=True, + bias=0.5, + beta_init=0.3, + inside_outside=False, + use_appearance_embedding=False, + position_encoding_max_degree=6, + use_numerical_gradients=True, + base_res=64, + max_res=4096, + log2_hashmap_size=22, + hash_features_per_level=8, + hash_smoothstep=False, + use_position_encoding=False, + ), + background_model="mlp", + enable_progressive_hash_encoding=True, + enable_curvature_loss_schedule=True, + enable_numerical_gradients_schedule=True, + ), + ), + optimizers={ + "fields": { + "optimizer": AdamWOptimizerConfig(lr=1e-3, weight_decay=0.01, eps=1e-15), + # "scheduler": NeuSSchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=500000), + "scheduler": MultiStepWarmupSchedulerConfig(warm_up_end=5000, milestones=[300_000, 400_000], gamma=0.1), + }, + "field_background": { + "optimizer": AdamWOptimizerConfig(lr=1e-3, eps=1e-15), + "scheduler": MultiStepWarmupSchedulerConfig(warm_up_end=5000, milestones=[300_000, 400_000], gamma=0.1), + }, + }, + viewer=ViewerConfig(num_rays_per_chunk=1 << 15), + vis="viewer", +) def merge_methods(methods, method_descriptions, new_methods, new_descriptions, overwrite=True): """Merge new methods and descriptions into existing methods and descriptions. diff --git a/nerfstudio/engine/optimizers.py b/nerfstudio/engine/optimizers.py index b80b77bcb..458e035f7 100644 --- a/nerfstudio/engine/optimizers.py +++ b/nerfstudio/engine/optimizers.py @@ -60,6 +60,13 @@ class AdamOptimizerConfig(OptimizerConfig): weight_decay: float = 0 """The weight decay to use.""" +@dataclass +class AdamWOptimizerConfig(OptimizerConfig): + """Basic optimizer config with AdamW""" + + _target: Type = torch.optim.AdamW + weight_decay: float = 0 + """The weight decay to use.""" @dataclass class RAdamOptimizerConfig(OptimizerConfig): diff --git a/nerfstudio/engine/schedulers.py b/nerfstudio/engine/schedulers.py index 6d1e11b66..85e8690e0 100644 --- a/nerfstudio/engine/schedulers.py +++ b/nerfstudio/engine/schedulers.py @@ -17,7 +17,7 @@ from abc import abstractmethod from dataclasses import dataclass, field -from typing import Literal, Optional, Tuple, Type +from typing import Literal, Optional, Tuple, Type, Any, List import numpy as np from torch.optim import Optimizer, lr_scheduler @@ -172,3 +172,35 @@ class CosineDecayScheduler(Scheduler): scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=func) return scheduler + +@dataclass +class MultiStepWarmupSchedulerConfig(InstantiateConfig): + """Basic scheduler config with self-defined exponential decay schedule""" + + _target: Type = field(default_factory=lambda: MultiStepWarmupScheduler) + warm_up_end: int = 5000 + milestones: List[int] = field(default_factory=lambda: [300000, 400000, 500000]) + gamma: float = 0.33 + + def setup(self, optimizer=None, **kwargs) -> Any: + """Returns the instantiated object using the config.""" + return self._target( + optimizer, + self.warm_up_end, + self.milestones, + self.gamma + ) + +class MultiStepWarmupScheduler(lr_scheduler.LambdaLR): + """Starts with a flat lr schedule until it reaches N epochs then applies a given scheduler""" + + def __init__(self, optimizer, warm_up_end, milestones, gamma) -> None: + def func(step): + if step < warm_up_end: + learning_factor = step / warm_up_end + else: + index = np.searchsorted(milestones, step, side='left') + learning_factor = gamma ** index + return learning_factor + + super().__init__(optimizer, lr_lambda=func) diff --git a/nerfstudio/models/neuralangelo.py b/nerfstudio/models/neuralangelo.py new file mode 100644 index 000000000..5039cb1d1 --- /dev/null +++ b/nerfstudio/models/neuralangelo.py @@ -0,0 +1,180 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Implementation of Neuralangelo model. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Dict, List, Tuple, Type + +import numpy as np +import torch +from torch.nn import Parameter + +from nerfstudio.cameras.rays import RayBundle +from nerfstudio.engine.callbacks import ( + TrainingCallback, + TrainingCallbackAttributes, + TrainingCallbackLocation, +) +from nerfstudio.field_components.field_heads import FieldHeadNames +from nerfstudio.fields.density_fields import HashMLPDensityField +from nerfstudio.model_components.losses import interlevel_loss +from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler +from nerfstudio.models.neus import NeuSModel, NeuSModelConfig +from nerfstudio.utils import colormaps + + +@dataclass +class NeuralangeloModelConfig(NeuSModelConfig): + """Neuralangelo Model Config""" + + _target: Type = field(default_factory=lambda: NeuralangeloModel) + #TODO move to base model config since it can be used in all models + enable_progressive_hash_encoding: bool = True + """whether to use progressive hash encoding""" + enable_numerical_gradients_schedule: bool = True + """whether to use numerical gradients delta schedule""" + enable_curvature_loss_schedule: bool = True + """whether to use curvature loss weight schedule""" + curvature_loss_multi: float = 5e-4 + """curvature loss weight""" + curvature_loss_warmup_steps: int = 5000 + """curvature loss warmup steps""" + level_init: int = 4 + """initial level of multi-resolution hash encoding""" + steps_per_level: int = 5000 + """steps per level of multi-resolution hash encoding""" + +class NeuralangeloModel(NeuSModel): + """Neuralangelo model + + Args: + config: Neuralangelo configuration to instantiate model + """ + + config: NeuralangeloModelConfig + + def populate_modules(self): + """Set the fields and modules.""" + super().populate_modules() + self.curvature_loss_multi_factor = 1.0 + + def get_training_callbacks( + self, training_callback_attributes: TrainingCallbackAttributes + ) -> List[TrainingCallback]: + callbacks = super().get_training_callbacks(training_callback_attributes) + + # read the hash encoding parameters from field + level_init = self.config.level_init + # schedule the delta in numerical gradients computation + num_levels = self.field.num_levels + max_res = self.field.max_res + base_res = self.field.base_res + growth_factor = self.field.growth_factor + + steps_per_level = self.config.steps_per_level + + + init_delta = 1. / base_res + end_delta = 1. / max_res + + # compute the delta based on level + if self.config.enable_numerical_gradients_schedule: + def set_delta(step): + delta = 1. / (base_res * growth_factor ** ( step / steps_per_level)) + delta = max(1. / max_res, delta) + self.field.set_numerical_gradients_delta(delta * 2.) # TODO because we divide 4 to normalize points to [0, 1] + + callbacks.append( + TrainingCallback( + where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION], + update_every_num_iters=1, + func=set_delta, + ) + ) + + # schedule the current level of multi-resolution hash encoding + if self.config.enable_progressive_hash_encoding: + def set_mask(step): + #TODO make this consistent with delta schedule + level = int(step / steps_per_level) + 1 + level = max(level, level_init) + self.field.update_mask(level) + + callbacks.append( + TrainingCallback( + where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION], + update_every_num_iters=1, + func=set_mask, + ) + ) + # schedule the curvature loss weight + # linear warmup for 5000 steps to 5e-4 and then decay as delta + if self.config.enable_curvature_loss_schedule: + def set_curvature_loss_mult_factor(step): + if step < self.config.curvature_loss_warmup_steps: + factor = step / self.config.curvature_loss_warmup_steps + else: + delta = 1. / (base_res * growth_factor ** ( (step - self.config.curvature_loss_warmup_steps) / steps_per_level)) + delta = max(1. / max_res, delta) + factor = delta / init_delta + + self.curvature_loss_multi_factor = factor + + callbacks.append( + TrainingCallback( + where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION], + update_every_num_iters=1, + func=set_curvature_loss_mult_factor, + ) + ) + + + #TODO switch to analytic gradients after delta is small enough? + + return callbacks + + def get_metrics_dict(self, outputs, batch) -> Dict: + metrics_dict = super().get_metrics_dict(outputs, batch) + + if self.training: + # training statics + metrics_dict["activated_encoding"] = self.field.hash_encoding_mask.mean().item() + metrics_dict["numerical_gradients_delta"] = self.field.numerical_gradients_delta + metrics_dict["curvature_loss_multi"] = self.curvature_loss_multi_factor * self.config.curvature_loss_multi + + return metrics_dict + + def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict: + loss_dict = super().get_loss_dict(outputs, batch, metrics_dict) + + # curvature loss + if self.training and self.config.curvature_loss_multi > 0.0: + delta = self.field.numerical_gradients_delta + centered_sdf = outputs['field_outputs'][FieldHeadNames.SDF] + sourounding_sdf = outputs['field_outputs']["sampled_sdf"] + + sourounding_sdf = sourounding_sdf.reshape(centered_sdf.shape[:2] + (3, 2)) + + # (a - b)/d - (b -c)/d = (a + c - 2b)/d + # ((a - b)/d - (b -c)/d)/d = (a + c - 2b)/(d*d) + curvature = (sourounding_sdf.sum(dim=-1) - 2 * centered_sdf) / (delta * delta) + loss_dict["curvature_loss"] = torch.abs(curvature).mean() * self.config.curvature_loss_multi * self.curvature_loss_multi_factor + + return loss_dict +