enable setting the scheduler

Author: weberpals

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.OneCycleLR]] = None
+    scheduler: Optional[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,16 +193,59 @@ def set_scheduler(self):
         Set the scheduler and scheduler args.
         The scheduler is not initialized yet as this is done in configure_optimizers in TimeNet.
         """
-        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,
-            }
-        )
+        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,
+                    }
+                )
+            elif self.scheduler.lower() == "steplr":
+                self.scheduler = torch.optim.lr_scheduler.StepLR
+                self.scheduler_args.update(
+                    {
+                        "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,
+                    }
+                )
+            elif self.scheduler.lower() == "reducelronplateau":
+                self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau
+                self.scheduler_args.update(
+                    {
+                        "mode": "min",
+                        "factor": 0.1,
+                        "patience": 10,
+                    }
+                )
+            elif self.scheduler.lower() == "cosineannealinglr":
+                self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR
+                self.scheduler_args.update(
+                    {
+                        "T_max": 50,
+                    }
+                )
+            else:
+                raise NotImplementedError(f"Scheduler {self.scheduler} is not supported.")
+        elif self.scheduler is None:
+            self.scheduler = torch.optim.lr_scheduler.ExponentialLR
+            self.scheduler_args.update(
+                {
+                    "gamma": 0.95,
+                }
+            )
 
     def set_lr_finder_args(self, dataset_size, num_batches):
         """

neuralprophet/forecaster.py

@@ -451,6 +451,7 @@ def __init__(
         accelerator: Optional[str] = None,
         trainer_config: dict = {},
         prediction_frequency: Optional[dict] = None,
+        scheduler: Optional[str] = "onecyclelr",
     ):
         self.config = locals()
         self.config.pop("self")
@@ -509,6 +510,7 @@ def __init__(
         self.config_train = configure.Train(
             quantiles=quantiles,
             learning_rate=learning_rate,
+            scheduler=scheduler,
             epochs=epochs,
             batch_size=batch_size,
             loss_func=loss_func,
@@ -921,6 +923,7 @@ def fit(
         continue_training: bool = False,
         num_workers: int = 0,
         deterministic: bool = False,
+        scheduler: Optional[str] = None,
     ):
         """Train, and potentially evaluate model.
 
@@ -986,6 +989,18 @@ def fit(
         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:
+                self.config_train.scheduler = 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."
+            )
+
         # Configuration
         if epochs is not None:
             self.config_train.epochs = epochs
@@ -2681,7 +2696,6 @@ def _init_train_loader(self, df, num_workers=0):
                 config_seasonality=self.config_seasonality,
             )
 
-        print("Changepoints:", self.config_trend.changepoints)
         df = _normalize(df=df, config_normalization=self.config_normalization)
         if not self.fitted:
             if self.config_trend.changepoints is not None:

neuralprophet/time_net.py

@@ -883,16 +883,23 @@ def configure_optimizers(self):
             for param_group in optimizer.param_groups:
                 param_group["initial_lr"] = (last_lr,)
 
-            lr_scheduler = lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(
-                optimizer, gamma=0.95, last_epoch=total_batches_processed - 1
-            )
-        else:
             lr_scheduler = self._scheduler(
                 optimizer,
-                max_lr=self.learning_rate,
-                total_steps=self.trainer.estimated_stepping_batches,
                 **self.config_train.scheduler_args,
             )
+        else:
+            if self._scheduler == torch.optim.lr_scheduler.OneCycleLR:
+                lr_scheduler = self._scheduler(
+                    optimizer,
+                    max_lr=self.learning_rate,
+                    total_steps=self.trainer.estimated_stepping_batches,
+                    **self.config_train.scheduler_args,
+                )
+            else:
+                lr_scheduler = self._scheduler(
+                    optimizer,
+                    **self.config_train.scheduler_args,
+                )
 
         return {"optimizer": optimizer, "lr_scheduler": lr_scheduler}