Add cosine annealing to LR scheduler options (#182)

Author: alan-cooney

sparse_autoencoder/train/sweep.py

@@ -138,15 +138,15 @@ def setup_optimizer(
     )
 
     lr_scheduler: LRScheduler | None = None
-
-    if hyperparameters["optimizer"]["lr_scheduler"] == "reduce_on_plateau":
-        lr_scheduler = ReduceLROnPlateau(optimizer=optim, patience=10)  # type: ignore
-
-    elif hyperparameters["optimizer"]["lr_scheduler"] == "cosine_annealing":
-        lr_scheduler = CosineAnnealingLR(
-            optimizer=optim,
-            T_max=10,
-        )
+    if "lr_scheduler" in hyperparameters["optimizer"]:
+        if hyperparameters["optimizer"]["lr_scheduler"] == "reduce_on_plateau":
+            lr_scheduler = ReduceLROnPlateau(optimizer=optim, patience=10)  # type: ignore
+
+        elif hyperparameters["optimizer"]["lr_scheduler"] == "cosine_annealing":
+            lr_scheduler = CosineAnnealingLR(
+                optimizer=optim,
+                T_max=10,
+            )
 
     return optim, lr_scheduler
 

sparse_autoencoder/train/sweep_config.py

@@ -152,7 +152,9 @@ class OptimizerHyperparameters(NestedParameter):
     Whether to use a fused implementation of the optimizer (may be faster on CUDA).
     """
 
-    lr_scheduler: Parameter[Literal["reduce_on_plateau"]] | None = None
+    lr_scheduler: Parameter[Literal["reduce_on_plateau", "cosine_annealing"]] | None = field(
+        default=None
+    )
     """Learning rate scheduler."""