LR warmup, angle sampling, huge interlevel loss

nerfstudio/configs/method_configs.py

@@ -46,7 +46,7 @@ from nerfstudio.data.dataparsers.phototourism_dataparser import (
     PhototourismDataParserConfig,
 )
 from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig
-from nerfstudio.engine.schedulers import SchedulerConfig,WarmupScheduler
+from nerfstudio.engine.schedulers import SchedulerConfig,WarmupSchedulerConfig
 from nerfstudio.engine.trainer import TrainerConfig
 from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind
 from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig
@@ -353,17 +353,17 @@ method_configs["dreamfusion"] = TrainerConfig(
         model=DreamFusionModelConfig(
             eval_num_rays_per_chunk=1 << 15,
             distortion_loss_mult=0.02,
-            interlevel_loss_mult=1.0,
+            interlevel_loss_mult=100,
             orientation_loss_mult=0.1,
             max_res=256,
             sphere_collider=True,
             initialize_density=False,
             random_background=True,
-            proposal_warmup=1000,
-            proposal_update_every=2,
+            proposal_warmup=500,
+            proposal_update_every=0,
             proposal_weights_anneal_max_num_iters=100,
-            start_lambertian_training=1000,
-            start_normals_training=500,
+            start_lambertian_training=500,
+            start_normals_training=1000,
             opacity_loss_mult=0.001,
             prompting_type="location_based",
             guidance_scale=100,
@@ -371,12 +371,12 @@ method_configs["dreamfusion"] = TrainerConfig(
     ),
     optimizers={
         "proposal_networks": {
-            "optimizer": AdamOptimizerConfig(lr=5e-2, eps=1e-15),
-            "scheduler": None,
+            "optimizer": AdamOptimizerConfig(lr=1e-8, eps=1e-15),
+            "scheduler": WarmupSchedulerConfig(lr_init=1e-8,lr_max=1e-3,lr_final=1e-6,warmup_steps=2000,max_steps=10000),
         },
         "fields": {
-            "optimizer": AdamOptimizerConfig(lr=5e-3, eps=1e-15),
-            "scheduler": WarmupScheduler(lr_init=1e-9,lr_max=1e-4,lr_final=1e-6,warmup_steps=3000,max_steps=10000),
+            "optimizer": AdamOptimizerConfig(lr=1e-8, eps=1e-15),
+            "scheduler": WarmupSchedulerConfig(lr_init=1e-8,lr_max=1e-3,lr_final=1e-6,warmup_steps=2000,max_steps=10000),
         },
     },
     viewer=ViewerConfig(num_rays_per_chunk=1 << 15),

nerfstudio/data/datamanagers/dreamfusion_datamanager.py

@@ -17,7 +17,7 @@ Data manager for dreamfusion
 """
 
 from __future__ import annotations
-
+import numpy as np
 from dataclasses import dataclass, field
 from typing import Dict, List, Tuple, Type, Union
 
@@ -167,9 +167,9 @@ class DreamFusionDataManagerConfig(VanillaDataManagerConfig):
     """Mean radius of camera orbit"""
     radius_std: float = 0.1
     """Std of radius of camera orbit"""
-    focal_range: Tuple[float, float] = (0.6, 1.2)
+    focal_range: Tuple[float, float] = (0.6, .1.2)
     """Range of focal length"""
-    vertical_rotation_range: Tuple[float, float] = (-80, 30)
+    vertical_rotation_range: Tuple[float, float] = (-90, 0)
     """Range of vertical rotation"""
     jitter_std: float = 0.05
     """Std of camera direction jitter, so we don't just point the cameras towards the center every time"""
@@ -245,6 +245,7 @@ class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-me
         #     return ray_bundle, {"initialization": False}
 
         # TODO below
+        ang_range = np.interp(step,[0,4000],[0,180])
         cameras, vertical_rotation, central_rotation = random_train_pose(
             self.config.train_images_per_batch,
             self.config.train_resolution,
@@ -255,6 +256,8 @@ class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-me
             vertical_rotation_range=self.config.vertical_rotation_range,
             jitter_std=self.config.jitter_std,
             center=self.config.center,
+            central_rotation_range = (-ang_range,ang_range)
+            
         )
         ray_bundle = cameras.generate_rays(torch.tensor(list(range(self.config.train_images_per_batch)))).flatten()
 

nerfstudio/engine/schedulers.py

@@ -72,6 +72,22 @@ class ExponentialDecaySchedule(lr_scheduler.LambdaLR):
 
         super().__init__(optimizer, lr_lambda=func)
 
+@dataclass
+class WarmupSchedulerConfig(InstantiateConfig):
+    """Basic scheduler config with self-defined warmup decay schedule"""
+
+    _target: Type = field(default_factory=lambda: WarmupScheduler)
+    lr_init:float=1e-8
+    lr_max:float=1e-3
+    lr_final:float=1e-5
+    warmup_steps:int = 1000
+    max_steps:int=10000
+    # TODO: somehow make this more generic. i dont like the idea of overriding the setup function
+    # but also not sure how to go about passing things into predefined torch objects.
+    def setup(self, optimizer=None, **kwargs) -> Any:
+        """Returns the instantiated object using the config."""
+        return self._target(optimizer,self.lr_init, self.lr_max, self.lr_final, self.warmup_steps, self.max_steps)
+
 class WarmupScheduler(lr_scheduler.LambdaLR):
     """Exponential learning rate decay function.
     See https://github.com/google-research/google-research/blob/
@@ -90,8 +106,7 @@ class WarmupScheduler(lr_scheduler.LambdaLR):
                 lr = np.interp(step,[0,warmup_steps],[lr_init,lr_max])
             else:
                 t = np.clip((step-warmup_steps)/max_steps,0,1)
-                lr = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t)
-            print("Lr",lr)
+                lr = np.exp(np.log(lr_max) * (1 - t) + np.log(lr_final) * t)
             return lr/lr_init#divide by lr_init to make the lr set to "lr"
         super().__init__(optimizer, lr_lambda=func)
 

nerfstudio/model_components/ray_samplers.py

@@ -263,7 +263,7 @@ class PDFSampler(Sampler):
         train_stratified: bool = True,
         single_jitter: bool = False,
         include_original: bool = True,
-        histogram_padding: float = 0.00,
+        histogram_padding: float = 0.005,
     ) -> None:
         super().__init__(num_samples=num_samples)
         self.train_stratified = train_stratified

nerfstudio/models/dreamfusion.py

@@ -114,7 +114,7 @@ class DreamFusionModelConfig(ModelConfig):
     """Slope of the annealing function for the proposal weights."""
     proposal_weights_anneal_max_num_iters: int = 500
     """Max num iterations for the annealing function."""
-    use_single_jitter: bool = True
+    use_single_jitter: bool = False
     """Whether use single jitter or not for the proposal networks."""
     interlevel_loss_mult: float = 1.0
     """Proposal loss multiplier."""
@@ -126,7 +126,7 @@ class DreamFusionModelConfig(ModelConfig):
     """start training with lambertian shading after this many iterations"""
     opacity_penalty: bool = True
     """enables penalty to encourage sparse weights (penalizing for uniform density along ray)"""
-    opacity_loss_mult: float = 1
+    opacity_loss_mult: float = 1e-3
     """scale for opacity penalty"""
     max_res: int = 256
     """Maximum resolution of the density field."""
@@ -341,8 +341,8 @@ class DreamFusionModel(Model):
 
     def get_outputs(self, ray_bundle: RayBundle):
         # uniform sampling
-        background_rgb = self.field.get_background_rgb(ray_bundle)
         ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns)
+        background_rgb = self.field.get_background_rgb(ray_bundle)
         field_outputs = self.field(ray_samples, compute_normals=True)
         weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY])
         weights_list.append(weights)
@@ -435,7 +435,7 @@ class DreamFusionModel(Model):
             assert weights.shape[-1] == 1
             if self.config.opacity_penalty:
                 outputs["opacity_loss"] = (
-                    torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * self.config.opacity_loss_mult
+                    torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01)
                 )
 
         return outputs