Merge pull request #1024 from mlcommons/scap3yvt-patch-2

Added new scheduler

Author: sarthakpati

.spelling/.spelling/expect.txt

@@ -747,3 +747,4 @@ srun
 tpu
 snfh
 netc
+cosineannealingwarmrestarts

GANDLF/configuration/scheduler_config.py

@@ -84,7 +84,9 @@ class SchedulerConfig(BaseModel):
     "reduce-on-plateau": ReduceOnPlateauConfig,
     "plateau": ReduceOnPlateauConfig,
     "reduceonplateau": ReduceOnPlateauConfig,
+    "cosineannealingwarmrestarts": CosineannealingConfig,
     "cosineannealing": CosineannealingConfig,
+    "cosineannealinglr": CosineannealingConfig,
     "warmupcosineschedule": WarmupcosinescheduleConfig,
     "wcs": WarmupcosinescheduleConfig,
 }

GANDLF/schedulers/__init__.py

@@ -6,7 +6,8 @@
     exp,
     step,
     reduce_on_plateau,
-    cosineannealing,
+    cosineannealingwarmrestarts,
+    cosineannealingLR,
 )
 
 from .wrap_monai import warmupcosineschedule
@@ -24,7 +25,9 @@
     "reduce-on-plateau": reduce_on_plateau,
     "plateau": reduce_on_plateau,
     "reduceonplateau": reduce_on_plateau,
-    "cosineannealing": cosineannealing,
+    "cosineannealing": cosineannealingwarmrestarts,
+    "cosineannealingwarmrestarts": cosineannealingwarmrestarts,
+    "cosineannealinglr": cosineannealingLR,
     "warmupcosineschedule": warmupcosineschedule,
     "wcs": warmupcosineschedule,
 }

GANDLF/schedulers/wrap_torch.py

@@ -5,6 +5,7 @@
     StepLR,
     ReduceLROnPlateau,
     CosineAnnealingWarmRestarts,
+    CosineAnnealingLR,
 )
 import math
 
@@ -169,14 +170,25 @@ def reduce_on_plateau(parameters):
     )
 
 
-def cosineannealing(parameters):
+def cosineannealingwarmrestarts(parameters):
     parameters["scheduler"]["T_0"] = parameters["scheduler"].get("T_0", 5)
     parameters["scheduler"]["T_mult"] = parameters["scheduler"].get("T_mult", 1)
-    parameters["scheduler"]["min_lr"] = parameters["scheduler"].get("min_lr", 0.001)
+    parameters["scheduler"]["eta_min"] = parameters["scheduler"].get("eta_min", 0.001)
 
     return CosineAnnealingWarmRestarts(
         parameters["optimizer_object"],
         T_0=parameters["scheduler"]["T_0"],
         T_mult=parameters["scheduler"]["T_mult"],
-        eta_min=parameters["scheduler"]["min_lr"],
+        eta_min=parameters["scheduler"]["eta_min"],
+    )
+
+
+def cosineannealingLR(parameters):
+    parameters["scheduler"]["T_max"] = parameters["scheduler"].get("T_max", 50)
+    parameters["scheduler"]["eta_min"] = parameters["scheduler"].get("eta_min", 0.001)
+
+    return CosineAnnealingLR(
+        parameters["optimizer_object"],
+        T_max=parameters["scheduler"]["T_max"],
+        eta_min=parameters["scheduler"]["eta_min"],
     )