add warmup scheduler

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
+from nerfstudio.engine.schedulers import SchedulerConfig,WarmupScheduler
 from nerfstudio.engine.trainer import TrainerConfig
 from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind
 from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig
@@ -376,7 +376,7 @@ method_configs["dreamfusion"] = TrainerConfig(
         },
         "fields": {
             "optimizer": AdamOptimizerConfig(lr=5e-3, eps=1e-15),
-            "scheduler": None,
+            "scheduler": WarmupScheduler(lr_init=1e-9,lr_max=1e-4,lr_final=1e-6,warmup_steps=3000,max_steps=10000),
         },
     },
     viewer=ViewerConfig(num_rays_per_chunk=1 << 15),

nerfstudio/engine/schedulers.py

@@ -72,6 +72,28 @@ class ExponentialDecaySchedule(lr_scheduler.LambdaLR):
 
         super().__init__(optimizer, lr_lambda=func)
 
+class WarmupScheduler(lr_scheduler.LambdaLR):
+    """Exponential learning rate decay function.
+    See https://github.com/google-research/google-research/blob/
+    fd2cea8cdd86b3ed2c640cbe5561707639e682f3/jaxnerf/nerf/utils.py#L360
+    for details.
+
+    Args:
+    
+    """
+
+    config: SchedulerConfig
+
+    def __init__(self, optimizer, lr_init, lr_max, lr_final, warmup_steps, max_steps ) -> None:
+        def func(step):
+            if step<warmup_steps:
+                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)
+            return lr/lr_init#divide by lr_init to make the lr set to "lr"
+        super().__init__(optimizer, lr_lambda=func)
 
 class DelayerScheduler(lr_scheduler.LambdaLR):
     """Starts with a flat lr schedule until it reaches N epochs then applies a given scheduler"""

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.01,
+        histogram_padding: float = 0.00,
     ) -> None:
         super().__init__(num_samples=num_samples)
         self.train_stratified = train_stratified