fix: num_loader_workers can be passed to historical_forecasts, only r…

…elevant for TorchForecastingModels

darts/models/forecasting/forecasting_model.py

@@ -328,10 +328,13 @@ def _predict_wrapper(
         num_samples: int,
         verbose: bool = False,
         predict_likelihood_parameters: bool = False,
+        num_loader_workers: int = 0,
     ) -> TimeSeries:
         kwargs = dict()
         if self.supports_likelihood_parameter_prediction:
             kwargs["predict_likelihood_parameters"] = predict_likelihood_parameters
+        if getattr(self, "trainer_params", False):
+            kwargs["num_loader_workers"] = num_loader_workers
         return self.predict(n, num_samples=num_samples, verbose=verbose, **kwargs)
 
     @property
@@ -586,6 +589,7 @@ def historical_forecasts(
         show_warnings: bool = True,
         predict_likelihood_parameters: bool = False,
         enable_optimization: bool = True,
+        num_loader_workers: int = 0,
     ) -> Union[
         TimeSeries, List[TimeSeries], Sequence[TimeSeries], Sequence[List[TimeSeries]]
     ]:
@@ -692,6 +696,9 @@ def historical_forecasts(
             Default: ``False``
         enable_optimization
             Whether to use the optimized version of historical_forecasts when supported and available.
+        num_loader_workers
+            Optionally, an integer specifying the ``num_workers`` to use in PyTorch ``DataLoader`` instances,
+            for the inference/prediction dataset loaders (if any).
 
         Returns
         -------
@@ -829,6 +836,7 @@ def retrain_func(
                 verbose=verbose,
                 show_warnings=show_warnings,
                 predict_likelihood_parameters=predict_likelihood_parameters,
+                num_loader_workers=num_loader_workers,
             )
 
         if len(series) == 1:
@@ -1019,6 +1027,7 @@ def retrain_func(
                     num_samples=num_samples,
                     verbose=verbose,
                     predict_likelihood_parameters=predict_likelihood_parameters,
+                    num_loader_workers=num_loader_workers,
                 )
                 if forecast_components is None:
                     forecast_components = forecast.columns
@@ -1076,6 +1085,7 @@ def backtest(
         reduction: Union[Callable[[np.ndarray], float], None] = np.mean,
         verbose: bool = False,
         show_warnings: bool = True,
+        num_loader_workers: int = 0,
     ) -> Union[float, List[float], Sequence[float], List[Sequence[float]]]:
         """Compute error values that the model would have produced when
         used on (potentially multiple) `series`.
@@ -1185,6 +1195,9 @@ def backtest(
             Whether to print progress.
         show_warnings
             Whether to show warnings related to parameters `start`, and `train_length`.
+        num_loader_workers
+            Optionally, an integer specifying the ``num_workers`` to use in PyTorch ``DataLoader`` instances,
+            for the inference/prediction dataset loaders (if any).
 
         Returns
         -------
@@ -1208,6 +1221,7 @@ def backtest(
                 last_points_only=last_points_only,
                 verbose=verbose,
                 show_warnings=show_warnings,
+                num_loader_workers=num_loader_workers,
             )
         else:
             forecasts = historical_forecasts
@@ -1261,6 +1275,7 @@ def gridsearch(
         verbose=False,
         n_jobs: int = 1,
         n_random_samples: Optional[Union[int, float]] = None,
+        num_loader_workers: int = 0,
     ) -> Tuple["ForecastingModel", Dict[str, Any], float]:
         """
         Find the best hyper-parameters among a given set using a grid search.
@@ -1374,6 +1389,9 @@ def gridsearch(
             must be between `0` and the total number of parameter combinations.
             If a float, `n_random_samples` is the ratio of parameter combinations selected from the full grid and must
             be between `0` and `1`. Defaults to `None`, for which random selection will be ignored.
+        num_loader_workers
+            Optionally, an integer specifying the ``num_workers`` to use in PyTorch ``DataLoader`` instances,
+            for the inference/prediction dataset loaders (if any).
 
         Returns
         -------
@@ -1457,6 +1475,7 @@ def _evaluate_combination(param_combination) -> float:
                     last_points_only=last_points_only,
                     verbose=verbose,
                     show_warnings=show_warnings,
+                    num_loader_workers=num_loader_workers,
                 )
             else:  # split mode
                 model._fit_wrapper(series, past_covariates, future_covariates)
@@ -1467,6 +1486,7 @@ def _evaluate_combination(param_combination) -> float:
                     future_covariates,
                     num_samples=1,
                     verbose=verbose,
+                    num_loader_workers=num_loader_workers,
                 )
                 error = metric(val_series, pred)
 
@@ -2131,6 +2151,7 @@ def predict(
         num_samples: int = 1,
         verbose: bool = False,
         predict_likelihood_parameters: bool = False,
+        num_loader_workers: int = 0,
     ) -> Union[TimeSeries, Sequence[TimeSeries]]:
         """Forecasts values for `n` time steps after the end of the series.
 
@@ -2169,6 +2190,9 @@ def predict(
             If set to `True`, the model predict the parameters of its Likelihood parameters instead of the target. Only
             supported for probabilistic models with a likelihood, `num_samples = 1` and `n<=output_chunk_length`.
             Default: ``False``
+        num_loader_workers
+            Optionally, an integer specifying the ``num_workers`` to use in PyTorch ``DataLoader`` instances,
+            for the inference/prediction dataset loaders (if any).
 
         Returns
         -------
@@ -2220,10 +2244,13 @@ def _predict_wrapper(
         num_samples: int,
         verbose: bool = False,
         predict_likelihood_parameters: bool = False,
+        num_loader_workers: int = 0,
     ) -> Union[TimeSeries, Sequence[TimeSeries]]:
         kwargs = dict()
         if self.supports_likelihood_parameter_prediction:
             kwargs["predict_likelihood_parameters"] = predict_likelihood_parameters
+        if getattr(self, "trainer_params", False):
+            kwargs["num_loader_workers"] = num_loader_workers
         return self.predict(
             n,
             series,
@@ -2465,10 +2492,13 @@ def _predict_wrapper(
         num_samples: int,
         verbose: bool = False,
         predict_likelihood_parameters: bool = False,
+        num_loader_workers: int = 0,
     ) -> TimeSeries:
         kwargs = dict()
         if self.supports_likelihood_parameter_prediction:
             kwargs["predict_likelihood_parameters"] = predict_likelihood_parameters
+        if getattr(self, "trainer_params", False):
+            kwargs["num_loader_workers"] = num_loader_workers
         return self.predict(
             n,
             future_covariates=future_covariates,
@@ -2682,10 +2712,13 @@ def _predict_wrapper(
         num_samples: int,
         verbose: bool = False,
         predict_likelihood_parameters: bool = False,
+        num_loader_workers: int = 0,
     ) -> TimeSeries:
         kwargs = dict()
         if self.supports_likelihood_parameter_prediction:
             kwargs["predict_likelihood_parameters"] = predict_likelihood_parameters
+        if getattr(self, "trainer_params", False):
+            kwargs["num_loader_workers"] = num_loader_workers
         return self.predict(
             n=n,
             series=series,

darts/models/forecasting/regression_model.py

@@ -1103,6 +1103,7 @@ def _optimized_historical_forecasts(
         verbose: bool = False,
         show_warnings: bool = True,
         predict_likelihood_parameters: bool = False,
+        num_loader_workers: int = 0,
     ) -> Union[
         TimeSeries, List[TimeSeries], Sequence[TimeSeries], Sequence[List[TimeSeries]]
     ]:

darts/models/forecasting/torch_forecasting_model.py

@@ -2037,6 +2037,7 @@ def _optimized_historical_forecasts(
         verbose: bool = False,
         show_warnings: bool = True,
         predict_likelihood_parameters: bool = False,
+        num_loader_workers: int = 0,
     ) -> Union[
         TimeSeries, List[TimeSeries], Sequence[TimeSeries], Sequence[List[TimeSeries]]
     ]:
@@ -2067,6 +2068,7 @@ def _optimized_historical_forecasts(
             show_warnings=show_warnings,
             predict_likelihood_parameters=predict_likelihood_parameters,
             verbose=verbose,
+            num_loader_workers=num_loader_workers,
         )
         return forecasts_list
 

darts/utils/historical_forecasts/optimized_historical_forecasts_regression.py

@@ -31,6 +31,7 @@ def _optimized_historical_forecasts_last_points_only(
     overlap_end: bool = False,
     show_warnings: bool = True,
     predict_likelihood_parameters: bool = False,
+    num_loader_workers: int = 0,
 ) -> Union[
     TimeSeries, List[TimeSeries], Sequence[TimeSeries], Sequence[List[TimeSeries]]
 ]:

darts/utils/historical_forecasts/optimized_historical_forecasts_torch.py

@@ -34,6 +34,7 @@ def _optimized_historical_forecasts(
     show_warnings: bool = True,
     predict_likelihood_parameters: bool = False,
     verbose: bool = False,
+    num_loader_workers: int = 0,
 ) -> Union[
     TimeSeries, List[TimeSeries], Sequence[TimeSeries], Sequence[List[TimeSeries]]
 ]:
@@ -99,6 +100,7 @@ def _optimized_historical_forecasts(
         future_covariates,
         num_samples=num_samples,
         predict_likelihood_parameters=predict_likelihood_parameters,
+        num_loader_workers=num_loader_workers,
     )
 
     dataset = model._build_inference_dataset(