Add lr scheduler interval config

Author: sorenmacbeth

src/pytorch_tabular/config/config.py

@@ -192,9 +192,9 @@ class DataConfig:
     )
 
     def __post_init__(self):
-        assert (
-            len(self.categorical_cols) + len(self.continuous_cols) + len(self.date_columns) > 0
-        ), "There should be at-least one feature defined in categorical, continuous, or date columns"
+        assert len(self.categorical_cols) + len(self.continuous_cols) + len(self.date_columns) > 0, (
+            "There should be at-least one feature defined in categorical, continuous, or date columns"
+        )
         _validate_choices(self)
         if os.name == "nt" and self.num_workers != 0:
             print("Windows does not support num_workers > 0. Setting num_workers to 0")
@@ -255,9 +255,9 @@ class InferredConfig:
 
     def __post_init__(self):
         if self.embedding_dims is not None:
-            assert all(
-                (isinstance(t, Iterable) and len(t) == 2) for t in self.embedding_dims
-            ), "embedding_dims must be a list of tuples (cardinality, embedding_dim)"
+            assert all((isinstance(t, Iterable) and len(t) == 2) for t in self.embedding_dims), (
+                "embedding_dims must be a list of tuples (cardinality, embedding_dim)"
+            )
             self.embedded_cat_dim = sum([t[1] for t in self.embedding_dims])
         else:
             self.embedded_cat_dim = 0
@@ -677,6 +677,9 @@ class OptimizerConfig:
         lr_scheduler_monitor_metric (Optional[str]): Used with ReduceLROnPlateau, where the plateau is
                 decided based on this metric
 
+        lr_scheduler_interval (Optional[str]): Interval at which to step the LR Scheduler, one of "epoch"
+                or "step". Defaults to `epoch`.
+
     """
 
     optimizer: str = field(
@@ -709,6 +712,11 @@ class OptimizerConfig:
         metadata={"help": "Used with ReduceLROnPlateau, where the plateau is decided based on this metric"},
     )
 
+    lr_scheduler_interval: Optional[str] = field(
+        default="epoch",
+        metadata={"help": "Interval at which to step the LR Scheduler, one of `epoch` or `step`. Defaults to `epoch`."},
+    )
+
     @staticmethod
     def read_from_yaml(filename: str = "config/optimizer_config.yml"):
         config = _read_yaml(filename)

src/pytorch_tabular/models/base_model.py

@@ -588,8 +588,11 @@ def configure_optimizers(self):
                 }
             return {
                 "optimizer": opt,
-                "lr_scheduler": self._lr_scheduler(opt, **self.hparams.lr_scheduler_params),
-                "monitor": self.hparams.lr_scheduler_monitor_metric,
+                "lr_scheduler": {
+                    "scheduler": self._lr_scheduler(opt, **self.hparams.lr_scheduler_params),
+                    "monitor": self.hparams.lr_scheduler_monitor_metric,
+                    "interval": self.hparams.lr_scheduler_interval,
+                },
             }
         else:
             return opt

src/pytorch_tabular/ssl_models/base_model.py

@@ -85,9 +85,9 @@ def __init__(
         self._setup_metrics()
 
     def _setup_encoder_decoder(self, encoder, encoder_config, decoder, decoder_config, inferred_config):
-        assert (encoder is not None) or (
-            encoder_config is not None
-        ), "Either encoder or encoder_config must be provided"
+        assert (encoder is not None) or (encoder_config is not None), (
+            "Either encoder or encoder_config must be provided"
+        )
         # assert (decoder is not None) or (decoder_config is not None),
         # "Either decoder or decoder_config must be provided"
         if encoder is not None:
@@ -181,7 +181,7 @@ def test_step(self, batch, batch_idx):
     def on_validation_epoch_end(self) -> None:
         if hasattr(self.hparams, "log_logits") and self.hparams.log_logits:
             warnings.warn(
-                "Logging Logits is disabled for SSL tasks. Set `log_logits` to False" " to turn off this warning"
+                "Logging Logits is disabled for SSL tasks. Set `log_logits` to False to turn off this warning"
             )
         super().on_validation_epoch_end()
 
@@ -219,8 +219,11 @@ def configure_optimizers(self):
                 }
             return {
                 "optimizer": opt,
-                "lr_scheduler": self._lr_scheduler(opt, **self.hparams.lr_scheduler_params),
-                "monitor": self.hparams.lr_scheduler_monitor_metric,
+                "lr_scheduler": {
+                    "scheduler": self._lr_scheduler(opt, **self.hparams.lr_scheduler_params),
+                    "monitor": self.hparams.lr_scheduler_monitor_metric,
+                    "interval": self.hparams.lr_scheduler_interval,
+                },
             }
         else:
             return opt