From 334acda31d564b73271e3f59e1b0c32c9bd7960c Mon Sep 17 00:00:00 2001
From: Justin <justin.g.kerr@gmail.com>
Date: Thu, 16 Feb 2023 20:04:06 +0000
Subject: [PATCH] add warmup scheduler

---
 nerfstudio/configs/method_configs.py        |  4 ++--
 nerfstudio/engine/schedulers.py             | 22 +++++++++++++++++++++
 nerfstudio/model_components/ray_samplers.py |  2 +-
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py
index d0c76262d..2ee48ca75 100644
--- a/nerfstudio/configs/method_configs.py
+++ b/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),
diff --git a/nerfstudio/engine/schedulers.py b/nerfstudio/engine/schedulers.py
index c545b9124..0f9df8da3 100644
--- a/nerfstudio/engine/schedulers.py
+++ b/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"""
diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py
index 08ee1fd95..49baefd98 100644
--- a/nerfstudio/model_components/ray_samplers.py
+++ b/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