Merge branch 'main' into nidhi_nn_losses

Author: phoeenniixx

docs/source/m_layer_v2.rst

@@ -47,3 +47,4 @@ See the detailed API documentation for the V2 base classes and specific model im
    models.samformer._samformer_v2.Samformer
    models.tide._tide_dsipts._tide_v2.TIDE
    models.timexer._timexer_v2.TimeXer
+   models.mlp._decodermlp_v2.DecoderMLP_v2

docs/source/pkg_v2.rst

@@ -99,3 +99,4 @@ See the detailed API documentation for the available V2 Package classes below:
    models.samformer._samformer_v2_pkg.Samformer_pkg_v2
    models.tide._tide_dsipts._tide_v2_pkg.TIDE_pkg_v2
    models.timexer._timexer_pkg_v2.TimeXer_pkg_v2
+   models.mlp._decodermlp_pkg_v2.DecoderMLP_pkg_v2

pyproject.toml

@@ -33,7 +33,7 @@ dependencies = [
   "scipy >=1.8,<2.0",
   "pandas >=1.3.0,<3.1.0",
   "scikit-learn >=1.2,<2.0",
-  "scikit-base <1.1.0",
+  "scikit-base <1.2.0",
 ]
 
 [project.optional-dependencies]

pytorch_forecasting/layers/__init__.py

@@ -19,6 +19,7 @@
     Encoder,
     EncoderLayer,
 )
+from pytorch_forecasting.layers._mlp import FullyConnectedModule
 from pytorch_forecasting.layers._normalization import RevIN
 from pytorch_forecasting.layers._output._flatten_head import (
     FlattenHead,
@@ -54,4 +55,5 @@
     "RevIN",
     "ResidualBlock",
     "embedding_cat_variables",
+    "FullyConnectedModule",
 ]

pytorch_forecasting/layers/_mlp/__init__.py

@@ -0,0 +1,7 @@
+"""
+Fully connected (MLP) layers.
+"""
+
+from pytorch_forecasting.layers._mlp._fully_connected import FullyConnectedModule
+
+__all__ = ["FullyConnectedModule"]

pytorch_forecasting/layers/_mlp/_fully_connected.py

@@ -0,0 +1,52 @@
+"""
+Fully connected (MLP) module.
+"""
+
+import torch
+from torch import nn
+
+
+class FullyConnectedModule(nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        output_size: int,
+        hidden_size: int,
+        n_hidden_layers: int,
+        activation_class: nn.ReLU,
+        dropout: float = None,
+        norm: bool = True,
+    ):
+        super().__init__()
+        self.input_size = input_size
+        self.output_size = output_size
+        self.hidden_size = hidden_size
+        self.n_hidden_layers = n_hidden_layers
+        self.activation_class = activation_class
+        self.dropout = dropout
+        self.norm = norm
+
+        # input layer
+        module_list = [nn.Linear(input_size, hidden_size), activation_class()]
+        if dropout is not None:
+            module_list.append(nn.Dropout(dropout))
+        if norm:
+            module_list.append(nn.LayerNorm(hidden_size))
+        # hidden layers
+        for _ in range(n_hidden_layers):
+            module_list.extend(
+                [nn.Linear(hidden_size, hidden_size), activation_class()]
+            )
+            if dropout is not None:
+                module_list.append(nn.Dropout(dropout))
+            if norm:
+                module_list.append(nn.LayerNorm(hidden_size))
+        # output layer
+        module_list.append(nn.Linear(hidden_size, output_size))
+
+        self.sequential = nn.Sequential(*module_list)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x of shape: batch_size x n_timesteps_in
+        # output of shape batch_size x n_timesteps_out
+        return self.sequential(x)

pytorch_forecasting/metrics/base_metrics/_base_metrics.py

@@ -703,18 +703,22 @@ def to_quantiles(self, y_pred: torch.Tensor, **kwargs) -> torch.Tensor:
         return self._metrics[0].to_quantiles(y_pred, **kwargs)
 
     def __add__(self, metric: LightningMetric):
+        new_metrics = list(self._metrics)
+        new_weights = list(self._weights)
         if isinstance(metric, self.__class__):
-            self._metrics.extend(metric._metrics)
-            self._weights.extend(metric._weights)
+            new_metrics.extend(metric._metrics)
+            new_weights.extend(metric._weights)
         else:
-            self._metrics.append(metric)
-            self._weights.append(1.0)
+            new_metrics.append(metric)
+            new_weights.append(1.0)
 
-        return self
+        result = CompositeMetric(metrics=new_metrics, weights=new_weights)
+        return result
 
     def __mul__(self, multiplier: float):
-        self._weights = [w * multiplier for w in self._weights]
-        return self
+        new_weights = [w * multiplier for w in self._weights]
+        result = CompositeMetric(metrics=list(self._metrics), weights=new_weights)
+        return result
 
     __rmul__ = __mul__
 

pytorch_forecasting/models/mlp/__init__.py

@@ -2,6 +2,14 @@
 
 from pytorch_forecasting.models.mlp._decodermlp import DecoderMLP
 from pytorch_forecasting.models.mlp._decodermlp_pkg import DecoderMLP_pkg
+from pytorch_forecasting.models.mlp._decodermlp_pkg_v2 import DecoderMLP_pkg_v2
+from pytorch_forecasting.models.mlp._decodermlp_v2 import DecoderMLP_v2
 from pytorch_forecasting.models.mlp.submodules import FullyConnectedModule
 
-__all__ = ["DecoderMLP", "DecoderMLP_pkg", "FullyConnectedModule"]
+__all__ = [
+    "DecoderMLP",
+    "DecoderMLP_pkg",
+    "DecoderMLP_v2",
+    "DecoderMLP_pkg_v2",
+    "FullyConnectedModule",
+]

pytorch_forecasting/models/mlp/_decodermlp_pkg_v2.py

@@ -0,0 +1,81 @@
+"""DecoderMLP v2 package container."""
+
+from pytorch_forecasting.base._base_pkg import Base_pkg
+
+
+class DecoderMLP_pkg_v2(Base_pkg):
+    """DecoderMLP v2 package container."""
+
+    _tags = {
+        "info:name": "DecoderMLP_v2",
+        "info:compute": 1,
+        "authors": ["jdb78", "echo-xiao"],
+        # TODO: the v2 datamodule supports categorical inputs but not
+        # categorical targets yet; add "categorical" to y_type once
+        # EncoderDecoderTimeSeriesDataModule supports categorical targets.
+        "info:y_type": ["numeric"],
+        "capability:exogenous": True,
+        "capability:multivariate": False,
+        "capability:pred_int": True,
+        "capability:flexible_history_length": True,
+        "capability:cold_start": True,
+    }
+
+    @classmethod
+    def get_cls(cls):
+        """Get model class."""
+        from pytorch_forecasting.models.mlp._decodermlp_v2 import DecoderMLP_v2
+
+        return DecoderMLP_v2
+
+    @classmethod
+    def get_datamodule_cls(cls):
+        """Get the underlying DataModule class."""
+        from pytorch_forecasting.data.data_module import (
+            EncoderDecoderTimeSeriesDataModule,
+        )
+
+        return EncoderDecoderTimeSeriesDataModule
+
+    @classmethod
+    def get_test_train_params(cls):
+        """Return testing parameter settings for the trainer.
+
+        Returns
+        -------
+        params : list of dict
+            Parameters to create testing instances of the class. Each dict is passed
+            as ``model_cfg`` to the package constructor; the ``"datamodule_cfg"`` key
+            is forwarded to the datamodule constructor.
+        """
+        from pytorch_forecasting.metrics import MAE, RMSE, SMAPE, QuantileLoss
+
+        params = [
+            {},
+            dict(
+                hidden_size=64, n_hidden_layers=2, dropout=0.1, norm=True, loss=RMSE()
+            ),
+            dict(
+                hidden_size=128,
+                n_hidden_layers=1,
+                activation_class="ReLU",
+                loss=SMAPE(),
+                logging_metrics=[MAE()],
+            ),
+            dict(hidden_size=32, n_hidden_layers=2, norm=False, loss=MAE()),
+            dict(hidden_size=64, n_hidden_layers=1, loss=QuantileLoss()),
+            dict(
+                optimizer="adamw",
+                lr_scheduler="cosine_annealing",
+                lr_scheduler_params={"T_max": 5},
+                loss=MAE(),
+            ),
+        ]
+
+        default_dm_cfg = {"max_encoder_length": 4, "max_prediction_length": 3}
+        for param in params:
+            dm_cfg = default_dm_cfg.copy()
+            dm_cfg.update(param.get("datamodule_cfg", {}))
+            param["datamodule_cfg"] = dm_cfg
+
+        return params

pytorch_forecasting/models/mlp/_decodermlp_v2.py

@@ -0,0 +1,171 @@
+"""Decoder-only MLP for pytorch-forecasting v2."""
+
+########################################################################################
+# Disclaimer: This implementation is based on the new v2 data pipeline and is
+# experimental, please use with care.
+########################################################################################
+
+import torch
+import torch.nn as nn
+from torch.optim import Optimizer
+
+from pytorch_forecasting.layers import FullyConnectedModule
+from pytorch_forecasting.metrics import QuantileLoss
+from pytorch_forecasting.models.base._base_model_v2 import BaseModel
+
+
+class DecoderMLP_v2(BaseModel):
+    """MLP on the decoder for pytorch-forecasting v2.
+
+    Predicts each future step purely from information known in the decoder
+    (future-known covariates and static features). It intentionally does not use the
+    encoder or target history -- it is a lightweight, covariate-driven baseline. The
+    original v1 ``DecoderMLP`` was authored by the pytorch-forecasting team.
+
+    Parameters
+    ----------
+    loss : nn.Module
+        Loss function for training (required).
+    hidden_size : int, default=300
+        Hidden layer width of the MLP.
+    n_hidden_layers : int, default=3
+        Number of hidden layers.
+    dropout : float, default=0.1
+        Dropout probability.
+    norm : bool, default=True
+        Whether to apply ``LayerNorm`` in the MLP.
+    activation_class : str, default="ReLU"
+        Name of a ``torch.nn`` activation class.
+    logging_metrics : Optional[list[nn.Module]], default=None
+        Metrics to log during training, validation, and testing.
+    optimizer : Optional[Union[Optimizer, str]], default="adam"
+        Optimizer to use for training.
+    optimizer_params : Optional[dict], default=None
+        Parameters for the optimizer.
+    lr_scheduler : Optional[str], default=None
+        Learning rate scheduler to use.
+    lr_scheduler_params : Optional[dict], default=None
+        Parameters for the learning rate scheduler.
+    metadata : Optional[dict], default=None
+        Metadata from ``EncoderDecoderTimeSeriesDataModule``.
+    """
+
+    @classmethod
+    def _pkg(cls):
+        """Package containing the model."""
+        from pytorch_forecasting.models.mlp._decodermlp_pkg_v2 import DecoderMLP_pkg_v2
+
+        return DecoderMLP_pkg_v2
+
+    def __init__(
+        self,
+        loss: nn.Module,
+        hidden_size: int = 300,
+        n_hidden_layers: int = 3,
+        dropout: float = 0.1,
+        norm: bool = True,
+        activation_class: str = "ReLU",
+        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 | None = None,
+    ):
+        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.hidden_size = hidden_size
+        self.n_hidden_layers = n_hidden_layers
+        self.dropout = dropout
+        self.norm = norm
+        self.activation_class = activation_class
+        self.metadata = metadata if metadata is not None else {}
+
+        self.save_hyperparameters(ignore=["loss", "logging_metrics", "metadata"])
+
+        # all dimensions are derived from metadata -- never hardcoded
+        self.decoder_cont_dim = self.metadata.get("decoder_cont", 0)
+        self.decoder_cat_dim = self.metadata.get("decoder_cat", 0)
+        self.static_cat_dim = self.metadata.get("static_categorical_features", 0)
+        self.static_cont_dim = self.metadata.get("static_continuous_features", 0)
+        self.prediction_length = self.metadata.get("max_prediction_length", 1)
+        self.target_dim = self.metadata.get("target", 1)
+
+        self._init_network()
+
+    def _init_network(self):
+        """Build the per-step MLP from the derived dimensions."""
+        self.input_size = (
+            self.decoder_cont_dim
+            + self.decoder_cat_dim
+            + self.static_cat_dim
+            + self.static_cont_dim
+        )
+
+        self.n_quantiles = None
+        if isinstance(self.loss, QuantileLoss):
+            self.n_quantiles = len(self.loss.quantiles)
+        self.output_size = self.n_quantiles if self.n_quantiles is not None else 1
+
+        self.mlp = FullyConnectedModule(
+            input_size=max(1, self.input_size),
+            output_size=self.output_size,
+            hidden_size=self.hidden_size,
+            n_hidden_layers=self.n_hidden_layers,
+            activation_class=getattr(nn, self.activation_class),
+            dropout=self.dropout,
+            norm=self.norm,
+        )
+
+    def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        """Forward pass: per-step MLP over decoder and static features.
+
+        Parameters
+        ----------
+        x : dict[str, torch.Tensor]
+            Input dictionary containing at least ``decoder_cont`` and, when the
+            corresponding metadata dimensions are non-zero, ``decoder_cat``,
+            ``static_continuous_features`` and ``static_categorical_features``.
+
+        Returns
+        -------
+        dict[str, torch.Tensor]
+            ``{"prediction": tensor}`` of shape
+            ``(batch_size, prediction_length, output_size)``.
+        """
+        decoder_cont = x["decoder_cont"]
+        batch_size = decoder_cont.shape[0]
+        pred_len = self.prediction_length
+        device = decoder_cont.device
+        dtype = decoder_cont.dtype
+
+        features = []
+        if self.decoder_cont_dim > 0:
+            features.append(x["decoder_cont"])
+        if self.decoder_cat_dim > 0:
+            features.append(x["decoder_cat"].to(dtype))
+        if self.static_cont_dim > 0:
+            features.append(x["static_continuous_features"].expand(-1, pred_len, -1))
+        if self.static_cat_dim > 0:
+            features.append(
+                x["static_categorical_features"].to(dtype).expand(-1, pred_len, -1)
+            )
+
+        if features:
+            network_input = torch.cat(features, dim=-1)
+        else:
+            network_input = torch.zeros(
+                batch_size, pred_len, 1, device=device, dtype=dtype
+            )
+
+        prediction = self.mlp(network_input.reshape(-1, self.mlp.input_size)).reshape(
+            batch_size, pred_len, self.output_size
+        )
+
+        return {"prediction": prediction}