add kwargs optimizer and scheduler to run_wrenformer() with defaults 'AdamW' and 'LambdaLR'

Author: janosh

examples/mat_bench/slurm_submit.py

@@ -18,7 +18,7 @@
 folds = list(range(5))
 checkpoint = None  # None | 'local' | 'wandb'
 lr = 3e-4
-model_name = f"wrenformer-{lr=:.0e}-{epochs=}-{n_attn_layers=}".replace("e-0", "e-")
+model_name = f"wrenformer-{lr=}-{epochs=}-{n_attn_layers=}"
 
 if "roost" in model_name.lower():
     # deploy Roost on all tasks

examples/mp_wbm/slurm_submit.py

@@ -13,17 +13,18 @@
 # %% write Python submission file and sbatch it
 epochs = 300
 n_attn_layers = 3
-embedding_aggregations = ("mean", "std", "min", "max")
+embedding_aggregations = ("mean",)
+optimizer = "AdamW"
+lr = 3e-4
+scheduler = ("CosineAnnealingLR", {"T_max": epochs})
 n_folds = 1
 data_path = f"{ROOT}/datasets/2022-06-09-mp+wbm.json.gz"
 target = "e_form"
 task_type = "regression"
 checkpoint = None  # None | 'local' | 'wandb'
-lr = 3e-4
-batch_size = 64
-model_name = f"wrenformer-robust-{lr=:.0e}-{epochs=}-{n_attn_layers=}".replace(
-    "e-0", "e-"
-)
+batch_size = 128
+model_name = f"wrenformer-robust-{epochs=}"
+
 
 os.makedirs(log_dir := f"{MODULE_DIR}/job-logs", exist_ok=True)
 timestamp = f"{datetime.now():%Y-%m-%d@%H-%M}"
@@ -52,7 +53,9 @@
     {epochs=},
     {n_attn_layers=},
     {checkpoint=},
+    {optimizer=},
     learning_rate={lr},
+    {scheduler=},
     {embedding_aggregations=},
     {batch_size=},
 )

examples/wrenformer.py

@@ -50,6 +50,8 @@ def run_wrenformer(
     checkpoint: Literal["local", "wandb"] | None = None,
     swa_start=0.7,
     run_params: dict[str, Any] = None,
+    optimizer: str | tuple[str, dict] = "AdamW",
+    scheduler: str | tuple[str, dict] = "LambdaLR",
     learning_rate: float = 3e-4,
     batch_size: int = 128,
     warmup_steps: int = 10,
@@ -82,9 +84,16 @@ def run_wrenformer(
             ```
         swa_start (float | None): When to start using stochastic weight averaging during training.
             Should be a float between 0 and 1. 0.7 means start SWA after 70% of epochs. Set to
-            None to disable SWA. Defaults to 0.7.
+            None to disable SWA. Defaults to 0.7. Proposed in https://arxiv.org/abs/1803.05407.
         run_params (dict[str, Any]): Additional parameters to merge into the run's dict of
             hyperparams. Will be logged to wandb. Can be anything really. Defaults to {}.
+        optimizer (str | tuple[str, dict]): Name of a torch.optim.Optimizer class like 'Adam',
+            'AdamW', 'SGD', etc. Can be a string or a string and dict with params to pass to the
+            class. Defaults to 'AdamW'.
+        scheduler (str | tuple[str, dict]): Name of a torch.optim.lr_scheduler class like
+            'LambdaLR', 'StepLR', 'CosineAnnealingLR', etc. Defaults to 'LambdaLR'. Can be a string
+            or a string and dict with params to pass to the class. E.g.
+            ('CosineAnnealingLR', {'T_max': n_epochs}).
         learning_rate (float): The optimizer's learning rate. Defaults to 3e-4.
         batch_size (int): The mini-batch size during training. Defaults to 128.
         warmup_steps (int): How many warmup steps the scheduler should do. Defaults to 10.
@@ -181,24 +190,46 @@ def run_wrenformer(
         embedding_aggregations=embedding_aggregations,
     )
     model.to(device)
-    optimizer = torch.optim.AdamW(params=model.parameters(), lr=learning_rate)
+    if isinstance(optimizer, str):
+        optimizer_name, optimizer_params = optimizer, None
+    elif isinstance(optimizer, (tuple, list)):
+        optimizer_name, optimizer_params = optimizer
+    else:
+        raise ValueError(f"Unknown {optimizer=}")
+    optimizer_cls = getattr(torch.optim, optimizer_name)
+    optimizer_instance = optimizer_cls(
+        params=model.parameters(), lr=learning_rate, **(optimizer_params or {})
+    )
 
     # This lambda goes up linearly until warmup_steps, then follows a power law decay.
     # Acts as a prefactor to the learning rate, i.e. actual_lr = lr_lambda(epoch) *
     # learning_rate.
-    scheduler = torch.optim.lr_scheduler.LambdaLR(
-        optimizer,
-        lambda epoch: min((epoch + 1) ** (-0.5), (epoch + 1) * warmup_steps ** (-1.5)),
-    )
+    if scheduler == "LambdaLR":
+        scheduler_name, scheduler_params = "LambdaLR", {
+            "lr_lambda": lambda epoch: min(
+                (epoch + 1) ** (-0.5), (epoch + 1) * warmup_steps ** (-1.5)
+            )
+        }
+    elif isinstance(scheduler, str):
+        scheduler_name, scheduler_params = scheduler, None
+    elif isinstance(scheduler, (tuple, list)):
+        scheduler_name, scheduler_params = scheduler
+    else:
+        raise ValueError(f"Unknown {scheduler=}")
+    scheduler_cls = getattr(torch.optim.lr_scheduler, scheduler_name)
+    scheduler_instance = scheduler_cls(optimizer_instance, **(scheduler_params or {}))
 
     if swa_start is not None:
         swa_model = AveragedModel(model)
-        # scheduler = CosineAnnealingLR(optimizer, T_max=epochs)
-        swa_scheduler = SWALR(optimizer, swa_lr=0.01)
+        swa_scheduler_instance = SWALR(optimizer_instance, swa_lr=0.01)
 
     run_params = {
         "epochs": epochs,
+        "optimizer": optimizer_name,
+        "optimizer_params": optimizer_params,
         "learning_rate": learning_rate,
+        "lr_scheduler": scheduler_name,
+        "scheduler_params": scheduler_params,
         "batch_size": batch_size,
         "n_attn_layers": n_attn_layers,
         "target": target_col,
@@ -232,7 +263,7 @@ def run_wrenformer(
         train_metrics = model.evaluate(
             train_loader,
             loss_dict,
-            optimizer,
+            optimizer_instance,
             normalizer_dict,
             action="train",
             verbose=verbose,
@@ -250,9 +281,9 @@ def run_wrenformer(
 
         if swa_start is not None and epoch > swa_start * epochs:
             swa_model.update_parameters(model)
-            swa_scheduler.step()
+            swa_scheduler_instance.step()
         else:
-            scheduler.step()
+            scheduler_instance.step()
 
         model.epoch += 1
 
@@ -293,9 +324,9 @@ def run_wrenformer(
     if checkpoint is not None:
         state_dict = {
             "model_state": inference_model.state_dict(),
-            "optimizer_state": optimizer.state_dict(),
+            "optimizer_state": optimizer_instance.state_dict(),
             "scheduler_state": (
-                scheduler if swa_start is None else swa_scheduler
+                scheduler_instance if swa_start is None else swa_scheduler_instance
             ).state_dict(),
             "loss_dict": loss_dict,
             "epoch": epochs,