robustify scheduler config

Author: ourownstory

neuralprophet/configure.py

@@ -94,7 +94,7 @@ class Train:
     optimizer: Union[str, Type[torch.optim.Optimizer]]
     quantiles: List[float] = field(default_factory=list)
     optimizer_args: dict = field(default_factory=dict)
-    scheduler: Optional[Type[torch.optim.lr_scheduler._LRScheduler]] = None
+    scheduler: Optional[Union[str, Type[torch.optim.lr_scheduler.LRScheduler]]] = None
     scheduler_args: dict = field(default_factory=dict)
     newer_samples_weight: float = 1.0
     newer_samples_start: float = 0.0
@@ -193,50 +193,48 @@ def set_scheduler(self):
         Set the scheduler and scheduler arg depending on the user selection.
         The scheduler is not initialized yet as this is done in configure_optimizers in TimeNet.
         """
-        self.scheduler_args.clear()
         if isinstance(self.scheduler, str):
             if self.scheduler.lower() == "onecyclelr":
                 self.scheduler = torch.optim.lr_scheduler.OneCycleLR
-                self.scheduler_args.update(
-                    {
-                        "pct_start": 0.3,
-                        "anneal_strategy": "cos",
-                        "div_factor": 10.0,
-                        "final_div_factor": 10.0,
-                        "three_phase": True,
-                    }
-                )
+                defaults = {
+                    "pct_start": 0.3,
+                    "anneal_strategy": "cos",
+                    "div_factor": 10.0,
+                    "final_div_factor": 10.0,
+                    "three_phase": True,
+                }
             elif self.scheduler.lower() == "steplr":
                 self.scheduler = torch.optim.lr_scheduler.StepLR
-                self.scheduler_args.update(
-                    {
-                        "step_size": 10,
-                        "gamma": 0.1,
-                    }
-                )
+                defaults = {
+                    "step_size": 10,
+                    "gamma": 0.1,
+                }
             elif self.scheduler.lower() == "exponentiallr":
                 self.scheduler = torch.optim.lr_scheduler.ExponentialLR
-                self.scheduler_args.update(
-                    {
-                        "gamma": 0.95,
-                    }
-                )
+                defaults = {
+                    "gamma": 0.95,
+                }
             elif self.scheduler.lower() == "cosineannealinglr":
                 self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR
-                self.scheduler_args.update(
-                    {
-                        "T_max": 50,
-                    }
-                )
+                defaults = {
+                    "T_max": 50,
+                }
             else:
-                raise NotImplementedError(f"Scheduler {self.scheduler} is not supported.")
+                raise NotImplementedError(
+                    f"Scheduler {self.scheduler} is not supported from string. Please pass the scheduler class."
+                )
+            if self.scheduler_args is not None:
+                defaults.update(self.scheduler_args)
+            self.scheduler_args = defaults
         elif self.scheduler is None:
             self.scheduler = torch.optim.lr_scheduler.ExponentialLR
-            self.scheduler_args.update(
-                {
-                    "gamma": 0.95,
-                }
-            )
+            self.scheduler_args = {
+                "gamma": 0.95,
+            }
+        else:  # if scheduler is a class
+            assert issubclass(
+                self.scheduler, torch.optim.lr_scheduler.LRScheduler
+            ), "Scheduler must be a subclass of torch.optim.lr_scheduler.LRScheduler"
 
     def set_lr_finder_args(self, dataset_size, num_batches):
         """

neuralprophet/forecaster.py

@@ -298,6 +298,7 @@ class NeuralProphet:
             >>> m = NeuralProphet(collect_metrics=["MSE", "MAE", "RMSE"])
             >>> # use custorm torchmetrics names
             >>> m = NeuralProphet(collect_metrics={"MAPE": "MeanAbsolutePercentageError", "MSLE": "MeanSquaredLogError",
+
         scheduler : str, torch.optim.lr_scheduler._LRScheduler
             Type of learning rate scheduler to use.
 
@@ -446,7 +447,8 @@ def __init__(
         batch_size: Optional[int] = None,
         loss_func: Union[str, torch.nn.modules.loss._Loss, Callable] = "SmoothL1Loss",
         optimizer: Union[str, Type[torch.optim.Optimizer]] = "AdamW",
-        scheduler: Optional[str] = "onecyclelr",
+        scheduler: Optional[Union[str, Type[torch.optim.lr_scheduler.LRScheduler]]] = "onecyclelr",
+        scheduler_args: Optional[dict] = None,
         newer_samples_weight: float = 2,
         newer_samples_start: float = 0.0,
         quantiles: List[float] = [],
@@ -521,6 +523,7 @@ def __init__(
             quantiles=quantiles,
             learning_rate=learning_rate,
             scheduler=scheduler,
+            scheduler_args=scheduler_args,
             epochs=epochs,
             batch_size=batch_size,
             loss_func=loss_func,
@@ -932,7 +935,8 @@ def fit(
         continue_training: bool = False,
         num_workers: int = 0,
         deterministic: bool = False,
-        scheduler: Optional[str] = None,
+        scheduler: Optional[Union[str, Type[torch.optim.lr_scheduler.LRScheduler]]] = None,
+        scheduler_args: Optional[dict] = None,
     ):
         """Train, and potentially evaluate model.
 
@@ -1002,20 +1006,30 @@ def fit(
                 "Model has been fitted already. If you want to continue training please set the flag continue_training."
             )
 
-        if continue_training and epochs is None:
-            raise ValueError("Continued training requires setting the number of epochs to train for.")
-
         if continue_training:
-            if scheduler is not None:
-                self.config_train.scheduler = scheduler
-            else:
+            if epochs is None:
+                raise ValueError("Continued training requires setting the number of epochs to train for.")
+
+            if continue_training and self.metrics_logger.checkpoint_path is None:
+                log.error("Continued training requires checkpointing in model to continue from last epoch.")
+
+            # if scheduler is not None:
+            #     log.warning(
+            #         "Scheduler can only be set in fit when continuing training. Please set the scheduler when initializing the model."
+            #     )
+
+            if scheduler is None:
+                log.warning(
+                    "No scheduler specified for continued training. Using a fallback scheduler for continued training."
+                )
                 self.config_train.scheduler = None
-            self.config_train.set_scheduler()
+                self.config_train.scheduler_args = None
+                self.config_train.set_scheduler()
 
-        if scheduler is not None and not continue_training:
-            log.warning(
-                "Scheduler can only be set in fit when continuing training. Please set the scheduler when initializing the model."
-            )
+        if scheduler is not None:
+            self.config_train.scheduler = scheduler
+            self.config_train.scheduler_args = scheduler_args
+            self.config_train.set_scheduler()
 
         # Configuration
         if epochs is not None:
@@ -1061,6 +1075,7 @@ def fit(
             log.info("When Global modeling with local normalization, metrics are displayed in normalized scale.")
 
         if minimal:
+            # overrides these settings:
             checkpointing = False
             self.metrics = False
             progress = None
@@ -1101,9 +1116,6 @@ def fit(
             or any(value != 1 for value in self.num_seasonalities_modelled_dict.values())
         )
 
-        if continue_training and self.metrics_logger.checkpoint_path is None:
-            log.error("Continued training requires checkpointing in model to continue from last epoch.")
-
         self.max_lags = df_utils.get_max_num_lags(
             n_lags=self.n_lags, config_lagged_regressors=self.config_lagged_regressors
         )