[ENH] Add Metrics support to ptf-v2 (#1960)

Fixes #1956, Fixes #1844
This PR tries to make `v2` compatible with `Metrics`. It also changes
the contract of tensors to match v1: `list` for multi-target and a
`tensor` for single-target

Stacks on #1965

Author: phoeenniixx

docs/source/tutorials/ptf_V2_example.ipynb

@@ -70,10 +70,65 @@ def __getitem__(self, idx: int) -> dict[str, Any]:
 
         Returns
         -------
-        x: dict[str, torch.Tensor]
-            A dictionary containing the processed data.
-        y: torch.Tensor
-            The target variable.
+        x : dict[str, torch.Tensor]
+            Dict containing processed inputs for the model, with the following keys:
+
+            * ``history_cont`` : torch.Tensor of shape
+                                    (context_length, n_history_cont_features)
+                Continuous features for the encoder (historical data).
+            * ``history_cat`` : torch.Tensor of shape
+                                    (context_length, n_history_cat_features)
+                Categorical features for the encoder (historical data).
+            * ``future_cont`` : torch.Tensor of shape
+                                    (prediction_length, n_future_cont_features)
+                Known continuous features for the decoder (future data).
+            * ``future_cat`` : torch.Tensor of shape
+                                    (prediction_length, n_future_cat_features)
+                Known categorical features for the decoder (future data).
+            * ``history_length`` : torch.Tensor of shape (1,)
+                Length of the encoder sequence.
+            * ``future_length`` : torch.Tensor of shape (1,)
+                Length of the decoder sequence.
+            * ``history_mask`` : torch.Tensor of shape (context_length,)
+                Boolean mask indicating valid encoder time points.
+            * ``future_mask`` : torch.Tensor of shape (prediction_length,)
+                Boolean mask indicating valid decoder time points.
+            * ``groups`` : torch.Tensor of shape (1,)
+                Group identifier for the time series instance.
+            * ``history_time_idx`` : torch.Tensor of shape (context_length,)
+                Time indices for the encoder sequence.
+            * ``future_time_idx`` : torch.Tensor of shape (prediction_length,)
+                Time indices for the decoder sequence.
+            * ``history_target`` : torch.Tensor of shape (context_length,)
+                Historical target values for the encoder sequence.
+            * ``future_target`` : torch.Tensor of shape (prediction_length,)
+                Target values for the decoder sequence.
+            * ``future_target_len`` : torch.Tensor of shape (1,)
+                Length of the decoder target sequence.
+
+            Optional fields, depending on dataset configuration:
+
+            * ``history_relative_time_idx`` : torch.Tensor of shape (context_length,),
+                                                optional
+                Relative time indices for the encoder sequence, present if
+                `add_relative_time_idx` is True.
+            * ``future_relative_time_idx`` : torch.Tensor of shape (prediction_length,),
+                                                optional
+                Relative time indices for the decoder sequence, present if
+                `add_relative_time_idx` is True.
+            * ``static_categorical_features`` : torch.Tensor of shape
+                                                (1, n_static_features), optional
+                Static categorical features if available.
+            * ``static_continuous_features`` : torch.Tensor of shape
+                                                (1, n_static_features), optional
+                Static continuous features if available.
+            * ``target_scale`` : torch.Tensor of shape (1,), optional
+                Scaling factor for the target values if provided by the dataset.
+
+        y : torch.Tensor or list of torch.Tensor
+            Target values for the decoder sequence.
+            If ``n_targets`` > 1, a list of tensors each of shape (prediction_length,)
+            is returned. Otherwise, a tensor of shape (prediction_length,) is returned.
         """
 
         series_idx, start_idx, context_length, prediction_length = self.windows[idx]
@@ -170,6 +225,10 @@ def __getitem__(self, idx: int) -> dict[str, Any]:
             x["target_scale"] = processed_data["target_scale"]
 
         y = processed_data["target"][future_indices]
+        if self.data_module.n_targets > 1:
+            y = [t.squeeze(-1) for t in torch.split(y, 1, dim=1)]
+        else:
+            y = y.squeeze(-1)
 
         return x, y
 
@@ -294,6 +353,7 @@ def __init__(
         self.window_stride = window_stride
 
         self.time_series_metadata = time_series_dataset.get_metadata()
+        self.n_targets = len(self.time_series_metadata["cols"]["y"])
 
         for idx, col in enumerate(self.time_series_metadata["cols"]["x"]):
             if self.time_series_metadata["col_type"].get(col) == "C":
@@ -774,8 +834,11 @@ def collate_fn(batch):
 
         Returns
         -------
-        tuple[dict[str, torch.Tensor], torch.Tensor]
+        tuple[dict[str, torch.Tensor], torch.Tensor or list of torch.Tensor]
             A tuple containing the collated data and the target variable.
+            If the dataset has multiple targets, a list of tensors each of shape
+            (batch_size, prediction_length,). Otherwise, a single tensor of shape
+            (batch_size, prediction_length).
         """
 
         x_batch = {
@@ -816,5 +879,13 @@ def collate_fn(batch):
                 [x["static_continuous_features"] for x, _ in batch]
             )
 
-        y_batch = torch.stack([y for _, y in batch])
+        if isinstance(batch[0][1], (list, tuple)):
+            num_targets = len(batch[0][1])
+            y_batch = []
+            for i in range(num_targets):
+                target_tensors = [sample_y[i] for _, sample_y in batch]
+                stacked_target = torch.stack(target_tensors)
+                y_batch.append(stacked_target)
+        else:
+            y_batch = torch.stack([y for _, y in batch])
         return x_batch, y_batch

docs/source/tutorials/tslib_v2_example.ipynb

@@ -127,7 +127,8 @@ def __init__(
         self.train_val_test_split = train_val_test_split
 
         warn(
-            "TimeSeries is part of an experimental rework of the "
+            "EncoderDecoderTimeSeriesDataModule is part of an experimental "
+            "rework of the "
             "pytorch-forecasting data layer, "
             "scheduled for release with v2.0.0. "
             "The API is not stable and may change without prior warning. "
@@ -151,6 +152,7 @@ def __init__(
         self._min_encoder_length = min_encoder_length or max_encoder_length
         self._categorical_encoders = _coerce_to_dict(categorical_encoders)
         self._scalers = _coerce_to_dict(scalers)
+        self.n_targets = len(self.time_series_metadata["cols"]["y"])
 
         self.categorical_indices = []
         self.continuous_indices = []
@@ -382,6 +384,13 @@ def __len__(self):
         def __getitem__(self, idx):
             """Retrieve a processed time series window for dataloader input.
 
+            Parameters
+            ----------
+            idx : int
+                Index of the window to retrieve from the dataset.
+
+            Returns
+            -------
             x : dict
                 Dictionary containing model inputs:
 
@@ -405,6 +414,8 @@ def __getitem__(self, idx):
                   Time indices for the encoder sequence.
                 * ``decoder_time_idx`` : tensor of shape (pred_length,)
                   Time indices for the decoder sequence.
+                * ``target_past`` : torch.Tensor of shape (enc_length,)
+                  Historical target values for the encoder sequence.
                 * ``target_scale`` : tensor of shape (1,)
                   Scaling factor for the target values.
                 * ``encoder_mask`` : tensor of shape (enc_length,)
@@ -420,8 +431,10 @@ def __getitem__(self, idx):
                 * ``static_continuous_features`` : tensor of shape (1, 0), optional
                   Placeholder for static continuous features (currently empty).
 
-            y : tensor of shape ``(pred_length, n_targets)``
+            y : torch.Tensor or list of torch.Tensor
                 Target values for the decoder sequence.
+                If ``n_targets`` > 1, a list of tensors each of shape (pred_length,)
+                is returned. Otherwise, a tensor of shape (pred_length,) is returned.
             """
             series_idx, start_idx, enc_length, pred_length = self.windows[idx]
             data = self.data_module._preprocess_data(series_idx)
@@ -547,8 +560,11 @@ def __getitem__(self, idx):
                     )
 
             y = data["target"][decoder_indices]
-            if y.ndim == 1:
-                y = y.unsqueeze(-1)
+
+            if self.data_module.n_targets > 1:
+                y = [t.squeeze(-1) for t in torch.split(y, 1, dim=1)]
+            else:
+                y = y.squeeze(-1)
 
             return x, y
 
@@ -730,5 +746,13 @@ def collate_fn(batch):
                 [x["static_continuous_features"] for x, _ in batch]
             )
 
-        y_batch = torch.stack([y for _, y in batch])
+        if isinstance(batch[0][1], (list, tuple)):
+            num_targets = len(batch[0][1])
+            y_batch = []
+            for i in range(num_targets):
+                target_tensors = [sample_y[i] for _, sample_y in batch]
+                stacked_target = torch.stack(target_tensors)
+                y_batch.append(stacked_target)
+        else:
+            y_batch = torch.stack([y for _, y in batch])
         return x_batch, y_batch

pytorch_forecasting/data/_tslib_data_module.py

@@ -14,39 +14,40 @@
 import torch.nn as nn
 from torch.optim import Optimizer
 
+from pytorch_forecasting.metrics import Metric
 from pytorch_forecasting.utils._classproperty import classproperty
 
 
 class BaseModel(LightningModule):
+    """Base model for time series forecasting.
+
+    Parameters
+    ----------
+    loss : Descendants of ``pytorch_forecasting.metrics.Metric`` class
+        Loss function to use for training.
+    logging_metrics : Optional[List[nn.Module]], optional
+        List of metrics to log during training, validation, and testing.
+    optimizer : Optional[Union[Optimizer, str]], optional
+        Optimizer to use for training.
+        Can be a string ("adam", "sgd") or an instance of `torch.optim.Optimizer`.
+    optimizer_params : Optional[Dict], optional
+        Parameters for the optimizer.
+    lr_scheduler : Optional[str], optional
+        Learning rate scheduler to use.
+        Supported values: "reduce_lr_on_plateau", "step_lr".
+    lr_scheduler_params : Optional[Dict], optional
+        Parameters for the learning rate scheduler.
+    """
+
     def __init__(
         self,
-        loss: nn.Module,
+        loss: Metric,
         logging_metrics: Optional[list[nn.Module]] = None,
         optimizer: Optional[Union[Optimizer, str]] = "adam",
         optimizer_params: Optional[dict] = None,
         lr_scheduler: Optional[str] = None,
         lr_scheduler_params: Optional[dict] = None,
     ):
-        """
-        Base model for time series forecasting.
-
-        Parameters
-        ----------
-        loss : nn.Module
-            Loss function to use for training.
-        logging_metrics : Optional[List[nn.Module]], optional
-            List of metrics to log during training, validation, and testing.
-        optimizer : Optional[Union[Optimizer, str]], optional
-            Optimizer to use for training.
-            Can be a string ("adam", "sgd") or an instance of `torch.optim.Optimizer`.
-        optimizer_params : Optional[Dict], optional
-            Parameters for the optimizer.
-        lr_scheduler : Optional[str], optional
-            Learning rate scheduler to use.
-            Supported values: "reduce_lr_on_plateau", "step_lr".
-        lr_scheduler_params : Optional[Dict], optional
-            Parameters for the learning rate scheduler.
-        """
         super().__init__()
         self.loss = loss
         self.logging_metrics = logging_metrics if logging_metrics is not None else []

pytorch_forecasting/data/data_module.py

@@ -9,6 +9,7 @@
 import torch.nn as nn
 from torch.optim import Optimizer
 
+from pytorch_forecasting.metrics import Metric
 from pytorch_forecasting.models.base._base_model_v2 import BaseModel
 
 
@@ -18,7 +19,7 @@ class TslibBaseModel(BaseModel):
 
     Parameters
     ----------
-    loss : nn.Module
+    loss : Descendants of ``pytorch_forecasting.metrics.Metric`` class
         Loss function to use for training.
     logging_metrics : Optional[list[nn.Module]], optional
         list of metrics to log during training, validation, and testing.
@@ -36,7 +37,7 @@ class TslibBaseModel(BaseModel):
 
     def __init__(
         self,
-        loss: nn.Module,
+        loss: Metric,
         logging_metrics: Optional[list[nn.Module]] = None,
         optimizer: Optional[Union[Optimizer, str]] = "adam",
         optimizer_params: Optional[dict] = None,

pytorch_forecasting/models/base/_base_model_v2.py

@@ -117,20 +117,20 @@ def get_test_train_params(cls):
 
         return [
             {
-                "loss": nn.MSELoss(),
+                # "loss": nn.MSELoss(),
                 "hidden_size": 32,
                 "use_revin": False,
             },
             {
-                "loss": nn.MSELoss(),
+                # "loss": nn.MSELoss(),
                 "hidden_size": 16,
                 "use_revin": True,
                 "out_channels": 1,
                 "persistence_weight": 0.0,
             },
-            # {
-            #     "loss": QuantileLoss(quantiles=[0.1, 0.5, 0.9]),
-            #     "hidden_size": 32,
-            #     "use_revin": False,
-            # },
+            {
+                "loss": QuantileLoss(quantiles=[0.1, 0.5, 0.9]),
+                "hidden_size": 32,
+                "use_revin": False,
+            },
         ]

pytorch_forecasting/models/base/_tslib_base_model_v2.py

@@ -109,7 +109,7 @@ def get_test_train_params(cls):
             `MyClass(**params)` or `MyClass(**params[i])` creates a valid test instance.
             `create_test_instance` uses the first (or only) dictionary in `params`
         """
-        import torch.nn as nn
+        from pytorch_forecasting.metrics import MAE, MAPE
 
         return [
             dict(
@@ -126,7 +126,7 @@ def get_test_train_params(cls):
                 n_add_dec=2,
                 dropout_rate=0.2,
                 data_loader_kwargs=dict(max_encoder_length=5, max_prediction_length=3),
-                loss=nn.MSELoss(),
+                loss=MAE(),
             ),
             dict(
                 hidden_size=64,
@@ -135,6 +135,6 @@ def get_test_train_params(cls):
                 n_add_dec=2,
                 dropout_rate=0.1,
                 data_loader_kwargs=dict(max_encoder_length=4, max_prediction_length=2),
-                loss=nn.PoissonNLLLoss(),
+                loss=MAPE(),
             ),
         ]

pytorch_forecasting/models/samformer/_samformer_v2_pkg.py

@@ -109,6 +109,8 @@ def get_test_train_params(cls):
             `MyClass(**params)` or `MyClass(**params[i])` creates a valid test instance.
             `create_test_instance` uses the first (or only) dictionary in `params`
         """
+        from pytorch_forecasting.metrics import QuantileLoss
+
         return [
             {},
             dict(
@@ -158,5 +160,6 @@ def get_test_train_params(cls):
                     context_length=16,
                     prediction_length=4,
                 ),
+                loss=QuantileLoss(quantiles=[0.1, 0.5, 0.9]),
             ),
         ]

pytorch_forecasting/models/tide/_tide_dsipts/_tide_v2_pkg.py

@@ -1,6 +1,5 @@
 """Automated tests based on the skbase test suite template."""
 
-from inspect import isclass
 import shutil
 
 import lightning.pytorch as pl
@@ -9,6 +8,7 @@
 import torch
 import torch.nn as nn
 
+from pytorch_forecasting.metrics import SMAPE
 from pytorch_forecasting.tests.test_all_estimators import (
     EstimatorFixtureGenerator,
     EstimatorPackageConfig,
@@ -62,7 +62,7 @@ def _integration(
         loss = kwargs["loss"]
         kwargs.pop("loss")
     else:
-        loss = nn.MSELoss()
+        loss = SMAPE()
 
     net = estimator_cls(
         metadata=metadata,