Configured optimizer and learning rate scheduling with a configurable warm-up length.

1. Previously, the linear learning rate in the schedulers (LinearLR, CosineAnnealingLR, and SequentialLR) was hard-coded to a single step. This has been replaced with the configurable `warmup_length`, enabling a more flexible and consistent setup p across all schedulers.

2. Defined a `warmup_length` parameter to allow consistent control over the linear warm-up phase.

Author: Hananeh Oliaei

src/electrai/configs/MP/config.yaml

@@ -30,8 +30,9 @@ epochs: 10
 nbatch: 1
 lr: 0.01
 weight_decay: 0.0
+warmup_length: 1
 
 # Weights and biases
-wandb_mode: ${env:WANDB_MODE}
+self.cfg.warmup_length: ${env:WANDB_MODE}
 wb_pname: ${env:WB_PNAME}
 entity: ${env:ENTITY}

src/electrai/lightning.py

@@ -56,12 +56,15 @@ def configure_optimizers(self):
         )
 
         linsch = torch.optim.lr_scheduler.LinearLR(
-            optimizer, start_factor=1e-5, end_factor=1, total_iters=1
+            optimizer,
+            start_factor=1e-5,
+            end_factor=1,
+            total_iters=self.cfg.warmup_length,
         )
         cossch = torch.optim.lr_scheduler.CosineAnnealingLR(
-            optimizer, T_max=int(self.cfg.epochs) - 1
+            optimizer, T_max=int(self.cfg.epochs) - self.cfg.warmup_length
         )
         scheduler = torch.optim.lr_scheduler.SequentialLR(
-            optimizer, [linsch, cossch], milestones=[1]
+            optimizer, [linsch, cossch], milestones=[self.cfg.warmup_length]
         )
         return [optimizer], [scheduler]