[devops] Upgrade to lightning 2.0

neuralprophet/configure.py

@@ -5,15 +5,15 @@
 import types
 from collections import OrderedDict
 from dataclasses import dataclass, field
-from typing import Callable, Iterable, List, Optional
+from typing import Callable, List, Optional
 from typing import OrderedDict as OrderedDictType
 from typing import Type, Union
 
 import numpy as np
 import pandas as pd
 import torch
 
-from neuralprophet import df_utils, np_types, utils, utils_torch
+from neuralprophet import df_utils, np_types, utils_torch
 from neuralprophet.custom_loss_metrics import PinballLoss
 from neuralprophet.hdays_utils import get_holidays_from_country
 
@@ -305,7 +305,7 @@
             log.error("Invalid growth for global_local mode '{}'. Set to 'global'".format(self.trend_global_local))
             self.trend_global_local = "global"
 
         if self.trend_local_reg < 0:
             log.error("Invalid  negative trend_local_reg '{}'. Set to False".format(self.trend_local_reg))
             self.trend_local_reg = False
 
@@ -354,13 +354,13 @@
             log.error("Invalid global_local mode '{}'. Set to 'global'".format(self.global_local))
             self.global_local = "global"
 
         self.periods = OrderedDict(
             {
                 "yearly": Season(
                     resolution=6,
                     period=365.25,
                     arg=self.yearly_arg,
                     global_local=(
                         self.yearly_global_local
                         if self.yearly_global_local in ["global", "local"]
                         else self.global_local
@@ -371,7 +371,7 @@
                     resolution=3,
                     period=7,
                     arg=self.weekly_arg,
                     global_local=(
                         self.weekly_global_local
                         if self.weekly_global_local in ["global", "local"]
                         else self.global_local
@@ -382,7 +382,7 @@
                     resolution=6,
                     period=1,
                     arg=self.daily_arg,
                     global_local=(
                         self.daily_global_local if self.daily_global_local in ["global", "local"] else self.global_local
                     ),
                     condition_name=None,
@@ -390,7 +390,7 @@
             }
         )
 
         if self.seasonality_local_reg < 0:
             log.error("Invalid  negative seasonality_local_reg '{}'. Set to False".format(self.seasonality_local_reg))
             self.seasonality_local_reg = False
 
@@ -410,7 +410,7 @@
             resolution=resolution,
             period=period,
             arg=arg,
             global_local=global_local if global_local in ["global", "local"] else self.global_local,
             condition_name=condition_name,
         )
 
@@ -486,7 +486,7 @@
     regressors: OrderedDict = field(init=False)  # contains RegressorConfig objects
 
     def __post_init__(self):
         self.regressors = None
 
 
 @dataclass

neuralprophet/custom_loss_metrics.py

@@ -33,15 +33,15 @@ def forward(self, outputs, target):
         """
         target = target.repeat(1, 1, len(self.quantiles))  # increase the quantile dimension of the targets
         differences = target - outputs
-        base_losses = self.loss_func(outputs, target)  # dimensions - [n_batch, n_forecasts, no. of quantiles]
-        positive_losses = (
-            torch.tensor(self.quantiles, device=target.device).unsqueeze(dim=0).unsqueeze(dim=0) * base_losses
+        base_losses = self.loss_func(outputs, target).float()  # dimensions - [n_batch, n_forecasts, no. of quantiles]
+        quantiles_tensor = (
+            torch.tensor(self.quantiles, device=target.device, dtype=torch.float32).unsqueeze(dim=0).unsqueeze(dim=0)
         )
-        negative_losses = (
-            1 - torch.tensor(self.quantiles, device=target.device).unsqueeze(dim=0).unsqueeze(dim=0)
-        ) * base_losses
+        positive_losses = quantiles_tensor * base_losses
+        negative_losses = (1 - quantiles_tensor) * base_losses
+        differences = differences.float()
         pinball_losses = torch.where(differences >= 0, positive_losses, negative_losses)
-        multiplier = torch.ones(size=(1, 1, len(self.quantiles)), device=target.device)
+        multiplier = torch.ones(size=(1, 1, len(self.quantiles)), device=target.device, dtype=torch.float32)
         multiplier[:, :, 0] = 2
         pinball_losses = multiplier * pinball_losses  # double the loss for the median quantile
         return pinball_losses

neuralprophet/forecaster.py

@@ -11,6 +11,7 @@
 import torch
 from matplotlib import pyplot
 from matplotlib.axes import Axes
+from pytorch_lightning.tuner.tuning import Tuner
 from torch.utils.data import DataLoader
 
 from neuralprophet import configure, df_utils, np_types, time_dataset, time_net, utils, utils_metrics
@@ -1104,7 +1105,7 @@
             # Only display the plot if the session is interactive, eg. do not show in github actions since it
             # causes an error in the Windows and MacOS environment
             if matplotlib.is_interactive():
                 fig
 
         self.fitted = True
         return metrics_df
@@ -2756,6 +2757,8 @@
         else:
             self.model = self._init_model()
 
+        self.model.train_loader = train_loader
+
         # Init the Trainer
         self.trainer, checkpoint_callback = utils.configure_trainer(
             config_train=self.config_train,
@@ -2780,8 +2783,9 @@
                 # Set parameters for the learning rate finder
                 self.config_train.set_lr_finder_args(dataset_size=dataset_size, num_batches=len(train_loader))
                 # Find suitable learning rate
-                lr_finder = self.trainer.tuner.lr_find(
-                    self.model,
+                tuner = Tuner(self.trainer)
+                lr_finder = tuner.lr_find(
+                    model=self.model,
                     train_dataloaders=train_loader,
                     val_dataloaders=val_loader,
                     **self.config_train.lr_finder_args,
@@ -2802,8 +2806,9 @@
                 # Set parameters for the learning rate finder
                 self.config_train.set_lr_finder_args(dataset_size=dataset_size, num_batches=len(train_loader))
                 # Find suitable learning rate
-                lr_finder = self.trainer.tuner.lr_find(
-                    self.model,
+                tuner = Tuner(self.trainer)
+                lr_finder = tuner.lr_find(
+                    model=self.model,
                     train_dataloaders=train_loader,
                     **self.config_train.lr_finder_args,
                 )
@@ -2831,7 +2836,6 @@
 
         if not metrics_enabled:
             return None
-
         # Return metrics collected in logger as dataframe
         metrics_df = pd.DataFrame(self.metrics_logger.history)
         return metrics_df

neuralprophet/logger.py

@@ -52,6 +52,7 @@ class ProgressBar(TQDMProgressBar):
     def __init__(self, *args, **kwargs):
         self.epochs = kwargs.pop("epochs")
         super().__init__(*args, **kwargs)
+        self.main_progress_bar = super().init_train_tqdm()
 
     def on_train_epoch_start(self, trainer: "pl.Trainer", *_) -> None:
         self.main_progress_bar.reset(self.epochs)

neuralprophet/time_net.py

@@ -795,7 +795,8 @@ def training_step(self, batch, batch_idx):
         scheduler.step()
 
         # Manually track the loss for the lr finder
-        self.trainer.fit_loop.running_loss.append(loss)
+        self.log("train_loss", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True)
+        self.log("reg_loss", reg_loss, on_step=False, on_epoch=True, prog_bar=True, logger=True)
 
         # Metrics
         if self.metrics_enabled:
@@ -983,6 +984,9 @@ def denormalize(self, ts):
             ts = scale_y * ts + shift_y
         return ts
 
+    def train_dataloader(self):
+        return self.train_loader
+
 
 class FlatNet(nn.Module):
     """

neuralprophet/utils.py

@@ -5,15 +5,14 @@
 import os
 import sys
 from collections import OrderedDict
-from typing import IO, TYPE_CHECKING, BinaryIO, Iterable, Optional, Union
+from typing import IO, TYPE_CHECKING, BinaryIO, Optional, Union
 
 import numpy as np
 import pandas as pd
 import pytorch_lightning as pl
 import torch
 
 from neuralprophet import utils_torch
-from neuralprophet.hdays_utils import get_country_holidays
 from neuralprophet.logger import ProgressBar
 
 if TYPE_CHECKING:
@@ -375,7 +374,7 @@
         seasonal_dims[name] = resolution
     return seasonal_dims
 
 
 def config_events_to_model_dims(config_events: Optional[ConfigEvents], config_country_holidays):
     """
     Convert user specified events configurations along with country specific
@@ -856,10 +855,6 @@
     """
     config = config.copy()
 
-    # Enable Learning rate finder if not learning rate provided
-    if config_train.learning_rate is None:
-        config["auto_lr_find"] = True
-
     # Set max number of epochs
     if hasattr(config_train, "epochs"):
         if config_train.epochs is not None: