add tests and adapt docstring

Author: weberpals

neuralprophet/configure.py

@@ -190,7 +190,7 @@ def set_optimizer(self):
 
     def set_scheduler(self):
         """
-        Set the scheduler and scheduler args.
+        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()
@@ -221,15 +221,6 @@ def set_scheduler(self):
                         "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(

neuralprophet/forecaster.py

@@ -301,6 +301,20 @@ 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.
+
+            Options
+                * (default) ``OneCycleLR``: One Cycle Learning Rate scheduler
+                * ``StepLR``: Step Learning Rate scheduler
+                * ``ExponentialLR``: Exponential Learning Rate scheduler
+                * ``CosineAnnealingLR``: Cosine Annealing Learning Rate scheduler
+
+            Examples
+            --------
+            >>> from neuralprophet import NeuralProphet
+            >>> # Step Learning Rate scheduler
+            >>> m = NeuralProphet(scheduler="StepLR")
 
         COMMENT
         Uncertainty Estimation
@@ -975,6 +989,13 @@ def fit(
                 Note: using multiple workers and therefore distributed training might significantly increase
                 the training time since each batch needs to be copied to each worker for each epoch. Keeping
                 all data on the main process might be faster for most datasets.
+            scheduler : str
+                Type of learning rate scheduler to use for continued training. If None, uses ExponentialLR as
+                default as specified in the model config.
+                Options
+                * ``StepLR``: Step Learning Rate scheduler
+                * ``ExponentialLR``: Exponential Learning Rate scheduler
+                * ``CosineAnnealingLR``: Cosine Annealing Learning Rate scheduler
 
         Returns
         -------
@@ -2796,7 +2817,8 @@ def _train(
             checkpoint_path = self.metrics_logger.checkpoint_path
             checkpoint = torch.load(checkpoint_path)
 
-            previous_epoch = self.model.current_epoch
+            checkpoint_epoch = checkpoint["epoch"] if "epoch" in checkpoint else 0
+            previous_epoch = max(self.model.current_epoch, checkpoint_epoch)
 
             # Set continue_training flag in model to update scheduler correctly
             self.model.continue_training = True

tests/test_utils.py

@@ -115,3 +115,34 @@ def test_continue_training():
     metrics = m.fit(df, checkpointing=True, freq="D")
     metrics2 = m.fit(df, freq="D", continue_training=True, epochs=ADDITIONAL_EPOCHS)
     assert metrics["Loss"].min() >= metrics2["Loss"].min()
+
+
+def test_continue_training_with_scheduler_selection():
+    df = pd.read_csv(PEYTON_FILE, nrows=NROWS)
+    m = NeuralProphet(
+        epochs=EPOCHS,
+        batch_size=BATCH_SIZE,
+        learning_rate=LR,
+        n_lags=6,
+        n_forecasts=3,
+        n_changepoints=0,
+    )
+    metrics = m.fit(df, checkpointing=True, freq="D")
+    # Continue training with StepLR
+    metrics2 = m.fit(df, freq="D", continue_training=True, epochs=ADDITIONAL_EPOCHS, scheduler="StepLR")
+    assert metrics["Loss"].min() >= metrics2["Loss"].min()
+
+
+def test_save_load_continue_training():
+    df = pd.read_csv(PEYTON_FILE, nrows=NROWS)
+    m = NeuralProphet(
+        epochs=EPOCHS,
+        n_lags=6,
+        n_forecasts=3,
+        n_changepoints=0,
+    )
+    metrics = m.fit(df, checkpointing=True, freq="D")
+    save(m, "test_model.pt")
+    m2 = load("test_model.pt")
+    metrics2 = m2.fit(df, continue_training=True, epochs=ADDITIONAL_EPOCHS, scheduler="StepLR")
+    assert metrics["Loss"].min() >= metrics2["Loss"].min()