rewrote TFTModel

darts/models/forecasting/block_rnn_model.py

@@ -11,11 +11,10 @@
 from darts.utils.likelihood_models import Likelihood
 from darts.logging import raise_if_not, get_logger
 from darts.utils.torch import random_method
-from darts.models.forecasting.ptl_model import (
-    PLTorchParametricProbabilisticForecastingModel as TorchParametricProbabilisticForecastingModel,
+from darts.models.forecasting.pl_forecasting_module import (
+    PLParametricProbabilisticForecastingModule as TorchParametricProbabilisticForecastingModel,
 )
 from darts.models.forecasting.torch_forecasting_model import (
-    # TorchParametricProbabilisticForecastingModel,
     PastCovariatesTorchModel,
 )
 

darts/models/forecasting/nbeats.py

@@ -13,11 +13,10 @@
 from darts.logging import get_logger, raise_log, raise_if_not
 from darts.utils.torch import random_method
 from darts.utils.likelihood_models import Likelihood
-from darts.models.forecasting.ptl_model import (
-    PLTorchParametricProbabilisticForecastingModel as TorchParametricProbabilisticForecastingModel,
+from darts.models.forecasting.pl_forecasting_module import (
+    PLParametricProbabilisticForecastingModule as TorchParametricProbabilisticForecastingModel,
 )
 from darts.models.forecasting.torch_forecasting_model import (
-    # TorchParametricProbabilisticForecastingModel,
     PastCovariatesTorchModel,
 )
 

darts/models/forecasting/ptl_model.py → darts/models/forecasting/pl_forecasting_module.py

@@ -1,5 +1,5 @@
 """
-This file contains abstract classes for deterministic and probabilistic pytorch-lightning modules
+This file contains abstract classes for deterministic and probabilistic PyTorch Lightning modules
 """
 
 import numpy as np
@@ -25,7 +25,7 @@
 
 
 # TODO: better names
-class PLTorchForecastingModel(pl.LightningModule, ABC):
+class PLForecastingModule(pl.LightningModule, ABC):
     @abstractmethod
     def __init__(
         self,
@@ -36,7 +36,7 @@ def __init__(
         lr_scheduler_kwargs: Optional[Dict] = None,
     ) -> None:
         """
-        Pytorch Lightning-based Forecasting Module.
+        PyTorch Lightning-based Forecasting Module.
 
         This class is meant to be inherited to create a new PyTorch Lightning-based forecasting module.
         When subclassing this class, please make sure to add the following methods with the given signatures:
@@ -67,7 +67,7 @@ def __init__(
             Optionally, some keyword arguments for the PyTorch optimizer.
         """
 
-        super(PLTorchForecastingModel, self).__init__()
+        super(PLForecastingModule, self).__init__()
 
         # Define the loss function
         self.criterion = loss_fn
@@ -98,7 +98,7 @@ def first_prediction_index(self) -> int:
 
     @abstractmethod
     def forward(self, *args, **kwargs) -> Any:
-        super(PLTorchForecastingModel, self).forward(*args, **kwargs)
+        super(PLForecastingModule, self).forward(*args, **kwargs)
 
     def training_step(self, train_batch, batch_idx) -> torch.Tensor:
         """performs the training step"""
@@ -267,7 +267,7 @@ def _sample_tiling(input_data_tuple, batch_sample_size):
         return tuple(tiled_input_data)
 
 
-class PLPastCovariatesTorchModel(PLTorchForecastingModel, ABC):
+class PLPastCovariatesModule(PLForecastingModule, ABC):
     def _produce_train_output(self, input_batch: Tuple):
         past_target, past_covariate = input_batch
         # Currently all our PastCovariates models require past target and covariates concatenated
@@ -372,14 +372,14 @@ def _produce_predict_output(self, x):
         return self.model(x)
 
 
-class PLFutureCovariatesTorchModel(PLTorchForecastingModel, ABC):
+class PLFutureCovariatesModule(PLForecastingModule, ABC):
     def _get_batch_prediction(
         self, n: int, input_batch: Tuple, roll_size: int
     ) -> Tensor:
         raise NotImplementedError("TBD: Darts doesn't contain such a model yet.")
 
 
-class PLDualCovariatesTorchModel(PLTorchForecastingModel, ABC):
+class PLDualCovariatesModule(PLForecastingModule, ABC):
     def _get_batch_prediction(
         self, n: int, input_batch: Tuple, roll_size: int
     ) -> Tensor:
@@ -388,7 +388,7 @@ def _get_batch_prediction(
         )
 
 
-class PLMixedCovariatesTorchModel(PLTorchForecastingModel, ABC):
+class PLMixedCovariatesModule(PLForecastingModule, ABC):
     def _get_batch_prediction(
         self, n: int, input_batch: Tuple, roll_size: int
     ) -> torch.Tensor:
@@ -539,15 +539,15 @@ def _get_batch_prediction(
         return batch_prediction
 
 
-class PLSplitCovariatesTorchModel(PLTorchForecastingModel, ABC):
+class PLSplitCovariatesModule(PLForecastingModule, ABC):
     def _get_batch_prediction(
         self, n: int, input_batch: Tuple, roll_size: int
     ) -> Tensor:
         raise NotImplementedError("TBD: Darts doesn't contain such a model yet.")
 
 
 # TODO: I think we could actually integrate probabilistic support already in the parent class and remove it from here?
-class PLTorchParametricProbabilisticForecastingModel(PLTorchForecastingModel, ABC):
+class PLParametricProbabilisticForecastingModule(PLForecastingModule, ABC):
     def __init__(self, likelihood: Optional[Likelihood] = None, **kwargs):
         """Pytorch Parametric Probabilistic Forecasting Model.
 

darts/models/forecasting/ptl_tft_model.py

@@ -11,7 +11,7 @@
 from torch import nn
 
 from darts import TimeSeries
-from darts.utils.torch import random_method
+
 from darts.logging import get_logger, raise_if_not, raise_if
 from darts.utils.likelihood_models import QuantileRegression, Likelihood
 from darts.utils.data import (
@@ -20,9 +20,9 @@
     MixedCovariatesTrainingDataset,
     MixedCovariatesInferenceDataset,
 )
-from darts.models.forecasting.ptl_model import (
-    PLTorchParametricProbabilisticForecastingModel,
-    PLMixedCovariatesTorchModel,
+from darts.models.forecasting.pl_forecasting_module import (
+    PLParametricProbabilisticForecastingModule,
+    PLMixedCovariatesModule,
 )
 from darts.models.forecasting.torch_forecasting_model import (
     MixedCovariatesTorchModel,
@@ -42,9 +42,7 @@
 ]
 
 
-class _TFTModule(
-    PLTorchParametricProbabilisticForecastingModel, PLMixedCovariatesTorchModel
-):
+class _TFTModule(PLParametricProbabilisticForecastingModule, PLMixedCovariatesModule):
     def __init__(
         self,
         output_dim: Tuple[int, int],
@@ -98,6 +96,8 @@ def __init__(
         likelihood
             The likelihood model to be used for probabilistic forecasts. By default the TFT uses
             a ``QuantileRegression`` likelihood.
+        kwargs
+            all parameters for Darts' :class:`PLForecastingModule`.
         """
 
         super(_TFTModule, self).__init__(likelihood=likelihood, **kwargs)

darts/models/forecasting/rnn_model.py

@@ -10,11 +10,10 @@
 from darts.timeseries import TimeSeries
 
 from darts.logging import raise_if_not, get_logger
-from darts.models.forecasting.ptl_model import (
-    PLTorchParametricProbabilisticForecastingModel as TorchParametricProbabilisticForecastingModel,
+from darts.models.forecasting.pl_forecasting_module import (
+    PLParametricProbabilisticForecastingModule as TorchParametricProbabilisticForecastingModel,
 )
 from darts.models.forecasting.torch_forecasting_model import (
-    # TorchParametricProbabilisticForecastingModel,
     DualCovariatesTorchModel,
 )
 from darts.utils.torch import random_method

darts/models/forecasting/tcn_model.py

@@ -15,11 +15,10 @@
 from darts.utils.likelihood_models import Likelihood
 
 from darts.logging import raise_if_not, get_logger
-from darts.models.forecasting.ptl_model import (
-    PLTorchParametricProbabilisticForecastingModel as TorchParametricProbabilisticForecastingModel,
+from darts.models.forecasting.pl_forecasting_module import (
+    PLParametricProbabilisticForecastingModule as TorchParametricProbabilisticForecastingModel,
 )
 from darts.models.forecasting.torch_forecasting_model import (
-    # TorchParametricProbabilisticForecastingModel,
     PastCovariatesTorchModel,
 )