adding neural angelo

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.

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):

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)

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
+