Feat/torchmetrics (#996)

* added droput and batch_norm similiar to nhits' implementation

* activation param for nhits

* fix

* support any pytorch activation function. NHiTs AvgPool1d support

* PR

* Update CHANGELOG.md

* fixed typo

* pytorch lightning did not like saving nn.modules

* first pass

* metrics now works with likelihood

* rename var

* made metrics a parameter. Added Tests

* torchmetrics is already a dependency

* changelog

* model now accepts torchmetrics and torchCollections

* torchmetric example in early stop

* handle no metrics

* made _calculate_metrics private

* Ray tune example

* Ray tune example

* changelog

* Update darts/models/forecasting/pl_forecasting_module.py

Co-authored-by: Dennis Bader <[email protected]>

* Update darts/models/forecasting/pl_forecasting_module.py

Co-authored-by: Dennis Bader <[email protected]>

* added torch_metrics to doc strings in torch based models

* added torch_metrics to doc strings in torch based models

* Update darts/models/forecasting/pl_forecasting_module.py

* black formatting

Co-authored-by: Greg DeVos <[email protected]>
Co-authored-by: Julien Herzen <[email protected]>
Co-authored-by: Dennis Bader <[email protected]>

Author: 4 people

CHANGELOG.md

@@ -15,6 +15,7 @@ Darts is still in an early development phase and we cannot always guarantee back
   by [Greg DeVos](https://github.com/gdevos010)
 - Implemented ["GLU Variants Improve Transformer"](https://arxiv.org/abs/2002.05202) for transformer based models (transformer and TFT). [#959](https://github.com/unit8co/darts/issues/959)
   by [Greg DeVos](https://github.com/gdevos010) 
+- Added support for torch metrics during training and validation. [#996](https://github.com/unit8co/darts/pull/996) by [Greg DeVos](https://github.com/gdevos010)
 
 ## [0.19.0](https://github.com/unit8co/darts/tree/0.19.0) (2022-04-13)
 ### For users of the library:

darts/models/forecasting/block_rnn_model.py

@@ -178,6 +178,9 @@ def __init__(
             PyTorch loss function used for training.
             This parameter will be ignored for probabilistic models if the ``likelihood`` parameter is specified.
             Default: ``torch.nn.MSELoss()``.
+        torch_metrics
+            A torch metric or a ``MetricCollection`` used for evaluation. A full list of available metrics can be found
+            at https://torchmetrics.readthedocs.io/en/latest/. Default: ``None``.
         likelihood
             One of Darts' :meth:`Likelihood <darts.utils.likelihood_models.Likelihood>` models to be used for
             probabilistic forecasts. Default: ``None``.

darts/models/forecasting/nbeats.py

@@ -598,6 +598,9 @@ def __init__(
             PyTorch loss function used for training.
             This parameter will be ignored for probabilistic models if the ``likelihood`` parameter is specified.
             Default: ``torch.nn.MSELoss()``.
+        torch_metrics
+            A torch metric or a ``MetricCollection`` used for evaluation. A full list of available metrics can be found
+            at https://torchmetrics.readthedocs.io/en/latest/. Default: ``None``.
         likelihood
             One of Darts' :meth:`Likelihood <darts.utils.likelihood_models.Likelihood>` models to be used for
             probabilistic forecasts. Default: ``None``.

darts/models/forecasting/nhits.py

@@ -534,6 +534,9 @@ def __init__(
             PyTorch loss function used for training.
             This parameter will be ignored for probabilistic models if the ``likelihood`` parameter is specified.
             Default: ``torch.nn.MSELoss()``.
+        torch_metrics
+            A torch metric or a ``MetricCollection`` used for evaluation. A full list of available metrics can be found
+            at https://torchmetrics.readthedocs.io/en/latest/. Default: ``None``.
         likelihood
             One of Darts' :meth:`Likelihood <darts.utils.likelihood_models.Likelihood>` models to be used for
             probabilistic forecasts. Default: ``None``.

darts/models/forecasting/pl_forecasting_module.py

@@ -3,11 +3,12 @@
 """
 
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Optional, Sequence, Tuple
+from typing import Any, Dict, Optional, Sequence, Tuple, Union
 
 import pytorch_lightning as pl
 import torch
 import torch.nn as nn
+import torchmetrics
 from joblib import Parallel, delayed
 
 from darts.logging import get_logger, raise_if, raise_log
@@ -29,10 +30,13 @@ def __init__(
         input_chunk_length: int,
         output_chunk_length: int,
         loss_fn: nn.modules.loss._Loss = nn.MSELoss(),
+        torch_metrics: Optional[
+            Union[torchmetrics.Metric, torchmetrics.MetricCollection]
+        ] = None,
         likelihood: Optional[Likelihood] = None,
         optimizer_cls: torch.optim.Optimizer = torch.optim.Adam,
         optimizer_kwargs: Optional[Dict] = None,
-        lr_scheduler_cls: torch.optim.lr_scheduler._LRScheduler = None,
+        lr_scheduler_cls: Optional[torch.optim.lr_scheduler._LRScheduler] = None,
         lr_scheduler_kwargs: Optional[Dict] = None,
     ) -> None:
         """
@@ -58,6 +62,9 @@ def __init__(
             PyTorch loss function used for training.
             This parameter will be ignored for probabilistic models if the ``likelihood`` parameter is specified.
             Default: ``torch.nn.MSELoss()``.
+        torch_metrics
+            A torch metric or a ``MetricCollection`` used for evaluation. A full list of available metrics can be found
+            at https://torchmetrics.readthedocs.io/en/latest/. Default: ``None``.
         likelihood
             One of Darts' :meth:`Likelihood <darts.utils.likelihood_models.Likelihood>` models to be used for
             probabilistic forecasts. Default: ``None``.
@@ -76,7 +83,8 @@ def __init__(
         super().__init__()
 
         # save hyper parameters for saving/loading
-        self.save_hyperparameters()
+        # do not save type nn.Module params
+        self.save_hyperparameters(ignore=["loss_fn", "torch_metrics"])
 
         raise_if(
             input_chunk_length is None or output_chunk_length is None,
@@ -100,6 +108,22 @@ def __init__(
             dict() if lr_scheduler_kwargs is None else lr_scheduler_kwargs
         )
 
+        if torch_metrics is None:
+            torch_metrics = torchmetrics.MetricCollection([])
+        elif isinstance(torch_metrics, torchmetrics.Metric):
+            torch_metrics = torchmetrics.MetricCollection([torch_metrics])
+        elif isinstance(torch_metrics, torchmetrics.MetricCollection):
+            pass
+        else:
+            raise_log(
+                AttributeError(
+                    "`torch_metrics` only accepts type torchmetrics.Metric or torchmetrics.MetricCollection"
+                ),
+                logger,
+            )
+        self.train_metrics = torch_metrics.clone(prefix="train_")
+        self.val_metrics = torch_metrics.clone(prefix="val_")
+
         # initialize prediction parameters
         self.pred_n: Optional[int] = None
         self.pred_num_samples: Optional[int] = None
@@ -126,6 +150,7 @@ def training_step(self, train_batch, batch_idx) -> torch.Tensor:
         ]  # By convention target is always the last element returned by datasets
         loss = self._compute_loss(output, target)
         self.log("train_loss", loss, batch_size=train_batch[0].shape[0], prog_bar=True)
+        self._calculate_metrics(output, target, self.train_metrics)
         return loss
 
     def validation_step(self, val_batch, batch_idx) -> torch.Tensor:
@@ -134,6 +159,7 @@ def validation_step(self, val_batch, batch_idx) -> torch.Tensor:
         target = val_batch[-1]
         loss = self._compute_loss(output, target)
         self.log("val_loss", loss, batch_size=val_batch[0].shape[0], prog_bar=True)
+        self._calculate_metrics(output, target, self.val_metrics)
         return loss
 
     def predict_step(
@@ -154,7 +180,7 @@ def predict_step(
         # number of individual series to be predicted in current batch
         num_series = input_data_tuple[0].shape[0]
 
-        # number of of times the input tensor should be tiled to produce predictions for multiple samples
+        # number of times the input tensor should be tiled to produce predictions for multiple samples
         # this variable is larger than 1 only if the batch_size is at least twice as large as the number
         # of individual time series being predicted in current batch (`num_series`)
         batch_sample_size = min(
@@ -226,12 +252,31 @@ def _compute_loss(self, output, target):
         if self.likelihood:
             return self.likelihood.compute_loss(output, target)
         else:
-            # If there's no likelihood, nr_params=1 and we need to squeeze out the
+            # If there's no likelihood, nr_params=1, and we need to squeeze out the
             # last dimension of model output, for properly computing the loss.
             return self.criterion(output.squeeze(dim=-1), target)
 
+    def _calculate_metrics(self, output, target, metrics):
+        if not len(metrics):
+            return
+
+        if self.likelihood:
+            _metric = metrics(target, self.likelihood.sample(output))
+        else:
+            # If there's no likelihood, nr_params=1, and we need to squeeze out the
+            # last dimension of model output, for properly computing the metric.
+            _metric = metrics(target, output.squeeze(dim=-1))
+
+        self.log_dict(
+            _metric,
+            on_epoch=True,
+            on_step=False,
+            logger=True,
+            prog_bar=True,
+        )
+
     def configure_optimizers(self):
-        """configures optimizers and learning rate schedulers for for model optimization."""
+        """configures optimizers and learning rate schedulers for model optimization."""
 
         # A utility function to create optimizer and lr scheduler from desired classes
         def _create_from_cls_and_kwargs(cls, kws):
@@ -365,7 +410,7 @@ def _get_batch_prediction(
         self, n: int, input_batch: Tuple, roll_size: int
     ) -> torch.Tensor:
         """
-        Feeds PastCovariatesTorchModel with input and output chunks of a PastCovariatesSequentialDataset to farecast
+        Feeds PastCovariatesTorchModel with input and output chunks of a PastCovariatesSequentialDataset to forecast
         the next ``n`` target values per target variable.
 
         Parameters:
@@ -416,7 +461,7 @@ def _get_batch_prediction(
                 batch_prediction[-1] = batch_prediction[-1][:, :roll_size, :]
 
             # ==========> PAST INPUT <==========
-            # roll over input series to contain latest target and covariate
+            # roll over input series to contain the latest target and covariate
             input_past = torch.roll(input_past, -roll_size, 1)
 
             # update target input to include next `roll_size` predictions
@@ -532,7 +577,7 @@ def _get_batch_prediction(
         self, n: int, input_batch: Tuple, roll_size: int
     ) -> torch.Tensor:
         """
-        Feeds MixedCovariatesModel with input and output chunks of a MixedCovariatesSequentialDataset to farecast
+        Feeds MixedCovariatesModel with input and output chunks of a MixedCovariatesSequentialDataset to forecast
         the next ``n`` target values per target variable.
 
         Parameters
@@ -598,7 +643,7 @@ def _get_batch_prediction(
                 batch_prediction[-1] = batch_prediction[-1][:, :roll_size, :]
 
             # ==========> PAST INPUT <==========
-            # roll over input series to contain latest target and covariate
+            # roll over input series to contain the latest target and covariate
             input_past = torch.roll(input_past, -roll_size, 1)
 
             # update target input to include next `roll_size` predictions

darts/models/forecasting/rnn_model.py

@@ -257,6 +257,9 @@ def __init__(
             PyTorch loss function used for training.
             This parameter will be ignored for probabilistic models if the ``likelihood`` parameter is specified.
             Default: ``torch.nn.MSELoss()``.
+        torch_metrics
+            A torch metric or a ``MetricCollection`` used for evaluation. A full list of available metrics can be found
+            at https://torchmetrics.readthedocs.io/en/latest/. Default: ``None``.
         likelihood
             One of Darts' :meth:`Likelihood <darts.utils.likelihood_models.Likelihood>` models to be used for
             probabilistic forecasts. Default: ``None``.

darts/models/forecasting/tcn_model.py

@@ -297,6 +297,9 @@ def __init__(
             PyTorch loss function used for training.
             This parameter will be ignored for probabilistic models if the ``likelihood`` parameter is specified.
             Default: ``torch.nn.MSELoss()``.
+        torch_metrics
+            A torch metric or a ``MetricCollection`` used for evaluation. A full list of available metrics can be found
+            at https://torchmetrics.readthedocs.io/en/latest/. Default: ``None``.
         likelihood
             One of Darts' :meth:`Likelihood <darts.utils.likelihood_models.Likelihood>` models to be used for
             probabilistic forecasts. Default: ``None``.

darts/models/forecasting/tft_model.py

@@ -629,6 +629,9 @@ def __init__(
             PyTorch loss function used for training. By default, the TFT model is probabilistic and uses a
             ``likelihood`` instead (``QuantileRegression``). To make the model deterministic, you can set the `
             `likelihood`` to None and give a ``loss_fn`` argument.
+        torch_metrics
+            A torch metric or a ``MetricCollection`` used for evaluation. A full list of available metrics can be found
+            at https://torchmetrics.readthedocs.io/en/latest/. Default: ``None``.
         likelihood
             The likelihood model to be used for probabilistic forecasts. By default, the TFT uses
             a ``QuantileRegression`` likelihood.

darts/models/forecasting/transformer_model.py

@@ -273,6 +273,9 @@ def __init__(
             PyTorch loss function used for training.
             This parameter will be ignored for probabilistic models if the ``likelihood`` parameter is specified.
             Default: ``torch.nn.MSELoss()``.
+        torch_metrics
+            A torch metric or a ``MetricCollection`` used for evaluation. A full list of available metrics can be found
+            at https://torchmetrics.readthedocs.io/en/latest/. Default: ``None``.
         likelihood
             One of Darts' :meth:`Likelihood <darts.utils.likelihood_models.Likelihood>` models to be used for
             probabilistic forecasts. Default: ``None``.