Repository: ourownstory/neural_prophet

Issue:
GPU support tensors
NeuralProphet does seem to raise an error when training on GPU's.

        running with CUDA does not work: RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0!

_Originally posted by @geneman in https://github.com/ourownstory/neural_prophet/issues/961#issuecomment-1331144742_
      


Discussion:
@geneman thanks for raising this issue, I will look into this.

Resolving patch:
diff --git a/neuralprophet/custom_loss_metrics.py b/neuralprophet/custom_loss_metrics.py
--- a/neuralprophet/custom_loss_metrics.py
+++ b/neuralprophet/custom_loss_metrics.py
@@ -34,10 +34,14 @@ 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).unsqueeze(dim=0).unsqueeze(dim=0) * base_losses
-        negative_losses = (1 - torch.tensor(self.quantiles).unsqueeze(dim=0).unsqueeze(dim=0)) * base_losses
+        positive_losses = (
+            torch.tensor(self.quantiles, device=target.device).unsqueeze(dim=0).unsqueeze(dim=0) * base_losses
+        )
+        negative_losses = (
+            1 - torch.tensor(self.quantiles, device=target.device).unsqueeze(dim=0).unsqueeze(dim=0)
+        ) * base_losses
         pinball_losses = torch.where(differences >= 0, positive_losses, negative_losses)
-        multiplier = torch.ones(size=(1, 1, len(self.quantiles)))
+        multiplier = torch.ones(size=(1, 1, len(self.quantiles)), device=target.device)
         multiplier[:, :, 0] = 2
         pinball_losses = multiplier * pinball_losses  # double the loss for the median quantile
         return pinball_losses
diff --git a/neuralprophet/time_net.py b/neuralprophet/time_net.py
--- a/neuralprophet/time_net.py
+++ b/neuralprophet/time_net.py
@@ -254,8 +254,11 @@ def __init__(
                     self.config_trend.changepoints = self.config_trend.changepoints_range * linear_t
                 else:
                     self.config_trend.changepoints = np.insert(self.config_trend.changepoints, 0, 0.0)
-                self.trend_changepoints_t = torch.tensor(
-                    self.config_trend.changepoints, requires_grad=False, dtype=torch.float
+                # Register in buffer so the tensor is moved to the correct device once initialized,
+                # https://pytorch-lightning.readthedocs.io/en/stable/starter/converting.html#remove-any-cuda-or-to-device-calls
+                self.register_buffer(
+                    "trend_changepoints_t",
+                    torch.tensor(self.config_trend.changepoints, requires_grad=False, dtype=torch.float),
                 )
 
                 # Trend Deltas parameters
@@ -513,7 +516,9 @@ def _compute_quantile_forecasts_from_diffs(self, diffs, predict_mode=False):
             if n_upper_quantiles > 0:  # check if upper quantiles exist
                 upper_quantile_diffs = diffs[:, :, quantiles_divider_index:]
                 if predict_mode:  # check for quantile crossing and correct them in predict mode
-                    upper_quantile_diffs[:, :, 0] = torch.max(torch.tensor(0), upper_quantile_diffs[:, :, 0])
+                    upper_quantile_diffs[:, :, 0] = torch.max(
+                        torch.tensor(0, device=self.device), upper_quantile_diffs[:, :, 0]
+                    )
                     for i in range(n_upper_quantiles - 1):
                         next_diff = upper_quantile_diffs[:, :, i + 1]
                         diff = upper_quantile_diffs[:, :, i]
@@ -525,7 +530,9 @@ def _compute_quantile_forecasts_from_diffs(self, diffs, predict_mode=False):
             if n_lower_quantiles > 0:  # check if lower quantiles exist
                 lower_quantile_diffs = diffs[:, :, 1:quantiles_divider_index]
                 if predict_mode:  # check for quantile crossing and correct them in predict mode
-                    lower_quantile_diffs[:, :, -1] = torch.max(torch.tensor(0), lower_quantile_diffs[:, :, -1])
+                    lower_quantile_diffs[:, :, -1] = torch.max(
+                        torch.tensor(0, device=self.device), lower_quantile_diffs[:, :, -1]
+                    )
                     for i in range(n_lower_quantiles - 1, 0, -1):
                         next_diff = lower_quantile_diffs[:, :, i - 1]
                         diff = lower_quantile_diffs[:, :, i]
@@ -769,7 +776,7 @@ def all_seasonalities(self, s, meta):
             torch.Tensor
                 Forecast component of dims (batch, n_forecasts)
         """
-        x = torch.zeros(size=(s[list(s.keys())[0]].shape[0], self.n_forecasts, len(self.quantiles)))
+        x = torch.zeros(size=(s[list(s.keys())[0]].shape[0], self.n_forecasts, len(self.quantiles)), device=self.device)
         for name, features in s.items():
             x = x + self.seasonality(features, name, meta)
         return x
@@ -924,17 +931,21 @@ def forward(self, inputs, meta=None):
             name_id_dummy = self.id_list[0]
             meta = OrderedDict()
             meta["df_name"] = [name_id_dummy for _ in range(inputs["time"].shape[0])]
-            meta = torch.tensor([self.id_dict[i] for i in meta["df_name"]])
+            meta = torch.tensor([self.id_dict[i] for i in meta["df_name"]], device=self.device)
         elif self.config_season is None:
             pass
         elif meta is None and self.config_season.global_local == "local":
             name_id_dummy = self.id_list[0]
             meta = OrderedDict()
             meta["df_name"] = [name_id_dummy for _ in range(inputs["time"].shape[0])]
-            meta = torch.tensor([self.id_dict[i] for i in meta["df_name"]])
+            meta = torch.tensor([self.id_dict[i] for i in meta["df_name"]], device=self.device)
 
-        additive_components = torch.zeros(size=(inputs["time"].shape[0], self.n_forecasts, len(self.quantiles)))
-        multiplicative_components = torch.zeros(size=(inputs["time"].shape[0], self.n_forecasts, len(self.quantiles)))
+        additive_components = torch.zeros(
+            size=(inputs["time"].shape[0], self.n_forecasts, len(self.quantiles)), device=self.device
+        )
+        multiplicative_components = torch.zeros(
+            size=(inputs["time"].shape[0], self.n_forecasts, len(self.quantiles)), device=self.device
+        )
 
         if "lags" in inputs:
             additive_components += self.auto_regression(lags=inputs["lags"])
@@ -1090,7 +1101,7 @@ def loss_func(self, inputs, predicted, targets):
             progress_in_epoch = 1 - ((steps_per_epoch * (self.current_epoch + 1) - self.global_step) / steps_per_epoch)
             loss, reg_loss = self._add_batch_regularizations(loss, self.current_epoch, progress_in_epoch)
         else:
-            reg_loss = torch.tensor(0.0)
+            reg_loss = torch.tensor(0.0, device=self.device)
         return loss, reg_loss
 
     def training_step(self, batch, batch_idx):
diff --git a/neuralprophet/utils.py b/neuralprophet/utils.py
--- a/neuralprophet/utils.py
+++ b/neuralprophet/utils.py
@@ -783,7 +783,7 @@ def configure_trainer(
             config["accelerator"] = accelerator
             config["devices"] = 1
 
-        if hasattr(config, "accelerator"):
+        if "accelerator" in config:
             log.info(f"Using accelerator {config['accelerator']} with {config['devices']} device(s).")
         else:
             log.info("No accelerator available. Using CPU for training.")


Test patch:
diff --git a/tests/test_integration.py b/tests/test_integration.py
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -1477,3 +1477,26 @@ def test_predict_raw():
     metrics = m.fit(df, freq="D")
     future = m.make_future_dataframe(df, periods=30, n_historic_predictions=100)
     forecast = m.predict(df=future, raw=True)
+
+
+def test_accelerator():
+    log.info("testing: accelerator in Lightning (if available)")
+    df = pd.read_csv(PEYTON_FILE, nrows=NROWS)
+    m = NeuralProphet(
+        n_forecasts=2,
+        n_lags=14,
+        num_hidden_layers=2,
+        d_hidden=32,
+        epochs=EPOCHS,
+        batch_size=BATCH_SIZE,
+        learning_rate=LR,
+        trend_reg=0.1,
+        quantiles=[0.1, 0.9],
+        accelerator="auto",
+    )
+    df["A"] = df["y"].rolling(7, min_periods=1).mean()
+    cols = [col for col in df.columns if col not in ["ds", "y"]]
+    m = m.add_lagged_regressor(names=cols)
+    m.highlight_nth_step_ahead_of_each_forecast(m.n_forecasts)
+    metrics_df = m.fit(df, freq="D")
+    forecast = m.predict(df)