diff --git a/pytorch_forecasting/models/__init__.py b/pytorch_forecasting/models/__init__.py index 6c88ca06f..5f98dd752 100644 --- a/pytorch_forecasting/models/__init__.py +++ b/pytorch_forecasting/models/__init__.py @@ -21,7 +21,11 @@ PatchTST_pkg_v2, PatchTST_v2, ) -from pytorch_forecasting.models.rnn import RecurrentNetwork +from pytorch_forecasting.models.rnn import ( + RecurrentNetwork, + RecurrentNetwork_pkg_v2, + RecurrentNetwork_v2, +) from pytorch_forecasting.models.scinet import SCINet_pkg_v2, SCINet_v2 from pytorch_forecasting.models.softs import SOFTS, SOFTS_pkg_v2 from pytorch_forecasting.models.temporal_fusion_transformer import ( @@ -41,6 +45,8 @@ "PatchTST_pkg_v2", "TemporalFusionTransformer", "RecurrentNetwork", + "RecurrentNetwork_v2", + "RecurrentNetwork_pkg_v2", "DeepAR", "BaseModel", "Baseline", diff --git a/pytorch_forecasting/models/rnn/__init__.py b/pytorch_forecasting/models/rnn/__init__.py index 23af65106..cceed05d4 100644 --- a/pytorch_forecasting/models/rnn/__init__.py +++ b/pytorch_forecasting/models/rnn/__init__.py @@ -2,5 +2,12 @@ from pytorch_forecasting.models.rnn._rnn import RecurrentNetwork from pytorch_forecasting.models.rnn._rnn_pkg import RecurrentNetwork_pkg +from pytorch_forecasting.models.rnn._rnn_pkg_v2 import RecurrentNetwork_pkg_v2 +from pytorch_forecasting.models.rnn._rnn_v2 import RecurrentNetwork_v2 -__all__ = ["RecurrentNetwork", "RecurrentNetwork_pkg"] +__all__ = [ + "RecurrentNetwork", + "RecurrentNetwork_pkg", + "RecurrentNetwork_v2", + "RecurrentNetwork_pkg_v2", +] diff --git a/pytorch_forecasting/models/rnn/_rnn_pkg_v2.py b/pytorch_forecasting/models/rnn/_rnn_pkg_v2.py new file mode 100644 index 000000000..c1da1fe70 --- /dev/null +++ b/pytorch_forecasting/models/rnn/_rnn_pkg_v2.py @@ -0,0 +1,65 @@ +"""RecurrentNetwork v2 package container.""" + +from typing import Any + +from pytorch_forecasting.base._base_pkg import Base_pkg + + +class RecurrentNetwork_pkg_v2(Base_pkg): + """RecurrentNetwork v2 package container.""" + + _tags: dict[str, Any] = { + "info:name": "RecurrentNetwork", + "info:compute": 1, + "info:y_type": ["numeric"], + "authors": ["jdb78", "Dev10-sys"], + "capability:exogenous": True, + "capability:multivariate": False, + "capability:pred_int": True, + "capability:flexible_history_length": True, + "capability:cold_start": False, + } + + @classmethod + def get_cls(cls): + """Get the model class.""" + from pytorch_forecasting.models.rnn._rnn_v2 import RecurrentNetwork_v2 + + return RecurrentNetwork_v2 + + @classmethod + def get_datamodule_cls(cls): + """Get the compatible DataModule class.""" + from pytorch_forecasting.data.data_module import ( + EncoderDecoderTimeSeriesDataModule, + ) + + return EncoderDecoderTimeSeriesDataModule + + @classmethod + def get_test_train_params(cls) -> list[dict[str, Any]]: + """Return testing parameter settings for trainer fixtures.""" + from pytorch_forecasting.metrics import MAE, QuantileLoss + + params: list[dict[str, Any]] = [ + {}, + {"cell_type": "GRU"}, + {"rnn_layers": 2, "dropout": 0.2}, + {"hidden_size": 8, "cell_type": "LSTM"}, + {"loss": MAE()}, + { + "loss": QuantileLoss(quantiles=[0.1, 0.5, 0.9]), + "hidden_size": 16, + }, + ] + + default_dm_cfg: dict[str, Any] = { + "max_encoder_length": 8, + "max_prediction_length": 3, + } + + for param in params: + current_dm_cfg = param.get("datamodule_cfg", {}) + param["datamodule_cfg"] = {**default_dm_cfg, **current_dm_cfg} + + return params diff --git a/pytorch_forecasting/models/rnn/_rnn_v2.py b/pytorch_forecasting/models/rnn/_rnn_v2.py new file mode 100644 index 000000000..389032a5d --- /dev/null +++ b/pytorch_forecasting/models/rnn/_rnn_v2.py @@ -0,0 +1,165 @@ +"""RecurrentNetwork (RNN / LSTM / GRU) model for PyTorch Forecasting v2.""" + +from typing import Any, Literal + +import torch +from torch import nn +from torch.optim import Optimizer + +from pytorch_forecasting.metrics import Metric +from pytorch_forecasting.models.base._base_model_v2 import BaseModel + + +class RecurrentNetwork_v2(BaseModel): + """RecurrentNetwork is a sequential time-series forecasting architecture + + built on standard PyTorch recurrent neural network layers (LSTM or GRU). + + Parameters + ---------- + loss : Metric + Loss function used during training. + cell_type : {"LSTM", "GRU"}, default="LSTM" + Type of recurrent cell to use. + hidden_size : int, default=10 + Hidden size of the recurrent layers. + rnn_layers : int, default=2 + Number of recurrent layers. + dropout : float, default=0.1 + Dropout rate applied between recurrent layers. + logging_metrics : list of nn.Module, optional + Metrics logged during training, validation, and testing. + optimizer : Optimizer or str, default="adam" + Optimizer used for training. + optimizer_params : dict, optional + Additional parameters for the optimizer. + lr_scheduler : str, optional + Learning rate scheduler name. + lr_scheduler_params : dict, optional + Parameters for the learning rate scheduler. + metadata : dict, optional + Metadata from the data module. Used to derive ``input_size``, + ``max_encoder_length``, and ``max_prediction_length``. + """ + + @classmethod + def _pkg(cls): + """Package container for the model.""" + from pytorch_forecasting.models.rnn._rnn_pkg_v2 import ( + RecurrentNetwork_pkg_v2, + ) + + return RecurrentNetwork_pkg_v2 + + def __init__( + self, + loss: Metric, + cell_type: Literal["LSTM", "GRU"] = "LSTM", + hidden_size: int = 10, + rnn_layers: int = 2, + dropout: float = 0.1, + logging_metrics: list[nn.Module] | None = None, + optimizer: Optimizer | str | None = "adam", + optimizer_params: dict | None = None, + lr_scheduler: str | None = None, + lr_scheduler_params: dict | None = None, + metadata: dict[str, Any] | None = None, + **kwargs: Any, + ): + super().__init__( + loss=loss, + logging_metrics=logging_metrics, + optimizer=optimizer, + optimizer_params=optimizer_params, + lr_scheduler=lr_scheduler, + lr_scheduler_params=lr_scheduler_params, + ) + self.save_hyperparameters(ignore=["loss", "logging_metrics", "metadata"]) + + self.cell_type = cell_type.upper() + if self.cell_type not in ["LSTM", "GRU"]: + raise ValueError( + f"Invalid cell_type: {cell_type}. Supported types are 'LSTM' and 'GRU'." + ) + + self.hidden_size = hidden_size + self.rnn_layers = rnn_layers + self.dropout = dropout + self.metadata = metadata or {} + + self.max_encoder_length = self.metadata.get("max_encoder_length", 10) + self.max_prediction_length = self.metadata.get("max_prediction_length", 1) + + # Calculate input feature dimensions from metadata + encoder_cont_dim = self.metadata.get("encoder_cont", 0) + target_dim = 1 + self.input_size = encoder_cont_dim + target_dim + + self.n_quantiles = 1 + quantiles = getattr(loss, "quantiles", None) + if quantiles is not None and hasattr(quantiles, "__len__"): + self.n_quantiles = len(quantiles) + + self.output_size = self.max_prediction_length * self.n_quantiles + + rnn_class = nn.LSTM if self.cell_type == "LSTM" else nn.GRU + self.rnn = rnn_class( + input_size=self.input_size, + hidden_size=self.hidden_size, + num_layers=self.rnn_layers, + dropout=self.dropout if self.rnn_layers > 1 else 0.0, + batch_first=True, + ) + + self.output_projector = nn.Linear(self.hidden_size, self.output_size) + + def _build_input_tensor(self, x: dict[str, torch.Tensor]) -> torch.Tensor: + """Build combined input tensor from continuous features and targets.""" + target_past = x.get("target_past") + if target_past is not None and target_past.ndim == 2: + target_past = target_past.unsqueeze(-1) + + encoder_cont = x.get("encoder_cont") + if encoder_cont is not None and encoder_cont.size(-1) > 0: + if target_past is not None: + return torch.cat([encoder_cont, target_past], dim=-1) + return encoder_cont + elif target_past is not None: + return target_past + else: + raise KeyError( + "Neither 'target_past' nor 'encoder_cont' found in input dict." + ) + + def forward( + self, + x: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + """Forward pass for RecurrentNetwork_v2. + + Parameters + ---------- + x : dict[str, torch.Tensor] + Dictionary containing input tensors from the DataModule. + + Returns + ------- + dict[str, torch.Tensor] + Dictionary containing predicted output tensor under key ``prediction``. + """ + input_tensor = self._build_input_tensor(x) + batch_size = input_tensor.size(0) + + rnn_out, _ = self.rnn(input_tensor) + last_hidden = rnn_out[:, -1, :] + + projected = self.output_projector(last_hidden) + + if self.n_quantiles > 1: + prediction = projected.view( + batch_size, self.max_prediction_length, self.n_quantiles + ) + else: + prediction = projected.unsqueeze(-1) + + return {"prediction": prediction}