Feat/torchmetrics

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/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, List, Optional, Sequence, Tuple
 
 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,12 @@ def __init__(
         input_chunk_length: int,
         output_chunk_length: int,
         loss_fn: nn.modules.loss._Loss = nn.MSELoss(),
+        torch_metrics: Optional[List[str]] = None,
+        metrics_params: Optional[List[Dict]] = 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 +61,11 @@ 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
+            List of torch metrics to be used for evaluation. A full list of available metrics can be found at
+            https://torchmetrics.readthedocs.io/en/latest/. Default: ``None``.
+        metrics_params
+            Dictionary of parameters to be passed to the metrics. Default: ``None``.
         likelihood
             One of Darts' :meth:`Likelihood <darts.utils.likelihood_models.Likelihood>` models to be used for
             probabilistic forecasts. Default: ``None``.
@@ -100,6 +108,23 @@ def __init__(
             dict() if lr_scheduler_kwargs is None else lr_scheduler_kwargs
         )
 
+        self.metrics = []
+        self.metrics_str = torch_metrics if torch_metrics else []
+        self.metrics_params = metrics_params if metrics_params else []
+
+        if metrics_params:
+            raise_if(
+                len(self.metrics_params) != len(self.metrics_str),
+                "Number of metrics parameters must be equal to number of metrics.",
+                logger,
+            )
+        # create empty dict for each metric
+        if self.metrics_str and metrics_params is None:
+            for _ in self.metrics_str:
+                self.metrics_params.append(dict())
+
+        self._setup_metrics()
+
         # initialize prediction parameters
         self.pred_n: Optional[int] = None
         self.pred_num_samples: Optional[int] = None
@@ -126,6 +151,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, tag="train")
         return loss
 
     def validation_step(self, val_batch, batch_idx) -> torch.Tensor:
@@ -134,6 +160,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, tag="valid")
         return loss
 
     def predict_step(
@@ -230,6 +257,45 @@ def _compute_loss(self, output, target):
             # last dimension of model output, for properly computing the loss.
             return self.criterion(output.squeeze(dim=-1), target)
 
+    def _setup_metrics(self):
+        if self.metrics_str:
+            self.metrics = []
+            task_module = torchmetrics.functional
+            for metric in self.metrics_str:
+                try:
+                    self.metrics.append(getattr(task_module, metric))
+                except AttributeError as e:
+                    raise_log(
+                        ValueError(
+                            f"{metric} is not a valid functional metric defined in the torchmetrics.functional module"
+                        )
+                    )
+                    raise e
+
+    def calculate_metrics(self, y, y_hat, tag):
+        metrics = []
+        for metric, metric_str, metric_params in zip(
+            self.metrics, self.metrics_str, self.metrics_params
+        ):
+            if self.likelihood:
+                _metric = metric(y_hat, self.likelihood.sample(y), **metric_params)
+            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 = metric(y_hat, y.squeeze(dim=-1), **metric_params)
+
+            metrics.append(_metric)
+
+            self.log(
+                f"{tag}_{metric_str}",
+                _metric,
+                on_epoch=True,
+                on_step=False,
+                logger=True,
+                prog_bar=True,
+            )
+        return metrics
+
     def configure_optimizers(self):
         """configures optimizers and learning rate schedulers for for model optimization."""
 

darts/tests/models/forecasting/test_torch_forecasting_model.py

@@ -21,13 +21,16 @@
     logger.warning("Torch not available. RNN tests will be skipped.")
     TORCH_AVAILABLE = False
 
-
 if TORCH_AVAILABLE:
 
     class TestTorchForecastingModel(DartsBaseTestClass):
         def setUp(self):
             self.temp_work_dir = tempfile.mkdtemp(prefix="darts")
 
+            times = pd.date_range("20130101", "20130410")
+            pd_series = pd.Series(range(100), index=times)
+            self.series = TimeSeries.from_series(pd_series)
+
         def tearDown(self):
             shutil.rmtree(self.temp_work_dir)
 
@@ -61,11 +64,8 @@ def test_suppress_automatic_save(self, patch_save_model):
                 save_checkpoints=False,
             )
 
-            times = pd.date_range("20130101", "20130410")
-            pd_series = pd.Series(range(100), index=times)
-            series = TimeSeries.from_series(pd_series)
-            model1.fit(series, epochs=1)
-            model2.fit(series, epochs=1)
+            model1.fit(self.series, epochs=1)
+            model2.fit(self.series, epochs=1)
 
             model1.predict(n=1)
             model2.predict(n=2)
@@ -101,12 +101,8 @@ def test_manual_save_and_load(self):
                 random_state=42,
             )
 
-            times = pd.date_range("20130101", "20130410")
-            pd_series = pd.Series(range(100), index=times)
-            series = TimeSeries.from_series(pd_series)
-
-            model_manual_save.fit(series, epochs=1)
-            model_auto_save.fit(series, epochs=1)
+            model_manual_save.fit(self.series, epochs=1)
+            model_auto_save.fit(self.series, epochs=1)
 
             model_dir = os.path.join(self.temp_work_dir)
 
@@ -215,10 +211,7 @@ def test_create_instance_existing_model_with_name_force_fit_with_reset(
             )
             # no exception is raised
 
-            times = pd.date_range("20130101", "20130410")
-            pd_series = pd.Series(range(100), index=times)
-            series = TimeSeries.from_series(pd_series)
-            model1.fit(series, epochs=1)
+            model1.fit(self.series, epochs=1)
 
             RNNModel(
                 12,
@@ -242,10 +235,7 @@ def test_train_from_0_n_epochs_20_no_fit_epochs(self):
                 12, "RNN", 10, 10, n_epochs=20, work_dir=self.temp_work_dir
             )
 
-            times = pd.date_range("20130101", "20130410")
-            pd_series = pd.Series(range(100), index=times)
-            series = TimeSeries.from_series(pd_series)
-            model1.fit(series)
+            model1.fit(self.series)
 
             self.assertEqual(20, model1.epochs_trained)
 
@@ -255,13 +245,10 @@ def test_train_from_20_n_epochs_40_no_fit_epochs(self):
                 12, "RNN", 10, 10, n_epochs=20, work_dir=self.temp_work_dir
             )
 
-            times = pd.date_range("20130101", "20130410")
-            pd_series = pd.Series(range(100), index=times)
-            series = TimeSeries.from_series(pd_series)
-            model1.fit(series)
+            model1.fit(self.series)
             self.assertEqual(20, model1.epochs_trained)
 
-            model1.fit(series)
+            model1.fit(self.series)
             self.assertEqual(20, model1.epochs_trained)
 
         # n_epochs = 20, fit|epochs=None, epochs_trained=10 - train for another 20 epochs
@@ -270,14 +257,11 @@ def test_train_from_10_n_epochs_20_no_fit_epochs(self):
                 12, "RNN", 10, 10, n_epochs=20, work_dir=self.temp_work_dir
             )
 
-            times = pd.date_range("20130101", "20130410")
-            pd_series = pd.Series(range(100), index=times)
-            series = TimeSeries.from_series(pd_series)
             # simulate the case that user interrupted training with Ctrl-C after 10 epochs
-            model1.fit(series, epochs=10)
+            model1.fit(self.series, epochs=10)
             self.assertEqual(10, model1.epochs_trained)
 
-            model1.fit(series)
+            model1.fit(self.series)
             self.assertEqual(20, model1.epochs_trained)
 
         # n_epochs = 20, fit|epochs=15, epochs_trained=10 - train for 15 epochs
@@ -286,20 +270,14 @@ def test_train_from_10_n_epochs_20_fit_15_epochs(self):
                 12, "RNN", 10, 10, n_epochs=20, work_dir=self.temp_work_dir
             )
 
-            times = pd.date_range("20130101", "20130410")
-            pd_series = pd.Series(range(100), index=times)
-            series = TimeSeries.from_series(pd_series)
             # simulate the case that user interrupted training with Ctrl-C after 10 epochs
-            model1.fit(series, epochs=10)
+            model1.fit(self.series, epochs=10)
             self.assertEqual(10, model1.epochs_trained)
 
-            model1.fit(series, epochs=15)
+            model1.fit(self.series, epochs=15)
             self.assertEqual(15, model1.epochs_trained)
 
         def test_optimizers(self):
-            times = pd.date_range("20130101", "20130410")
-            pd_series = pd.Series(range(100), index=times)
-            series = TimeSeries.from_series(pd_series)
 
             optimizers = [
                 (torch.optim.Adam, {"lr": 0.001}),
@@ -316,12 +294,9 @@ def test_optimizers(self):
                     optimizer_kwargs=optim_kwargs,
                 )
                 # should not raise an error
-                model.fit(series, epochs=1)
+                model.fit(self.series, epochs=1)
 
         def test_lr_schedulers(self):
-            times = pd.date_range("20130101", "20130410")
-            pd_series = pd.Series(range(100), index=times)
-            series = TimeSeries.from_series(pd_series)
 
             lr_schedulers = [
                 (torch.optim.lr_scheduler.StepLR, {"step_size": 10}),
@@ -342,7 +317,7 @@ def test_lr_schedulers(self):
                     lr_scheduler_kwargs=lr_scheduler_kwargs,
                 )
                 # should not raise an error
-                model.fit(series, epochs=1)
+                model.fit(self.series, epochs=1)
 
         def test_devices(self):
             torch_devices = [
@@ -373,3 +348,45 @@ def test_wrong_model_creation_params(self):
             # invalid params should raise an error
             with self.assertRaises(ValueError):
                 _ = RNNModel(12, "RNN", 10, 10, **invalid_kwarg)
+
+        def test_metrics(self):
+            torch_metrics = ["mean_squared_error", "mean_absolute_percentage_error"]
+            model = RNNModel(12, "RNN", 10, 10, n_epochs=1, torch_metrics=torch_metrics)
+            model.fit(self.series)
+
+        def test_metrics_w_params(self):
+            torch_metrics = ["mean_squared_error", "mean_absolute_percentage_error"]
+            metrics_params = [{}, {}]
+            model = RNNModel(
+                12,
+                "RNN",
+                10,
+                10,
+                n_epochs=1,
+                torch_metrics=torch_metrics,
+                metrics_params=metrics_params,
+            )
+            model.fit(self.series)
+
+        def test_invalid_metrics(self):
+            torch_metrics = ["invalid"]
+            with self.assertRaises(ValueError):
+                model = RNNModel(
+                    12, "RNN", 10, 10, n_epochs=1, torch_metrics=torch_metrics
+                )
+                model.fit(self.series)
+
+        def test_wrong_metrics_param_count(self):
+            torch_metrics = ["mean_squared_error", "mean_absolute_percentage_error"]
+            metrics_params = [{}]
+            with self.assertRaises(ValueError):
+                model = RNNModel(
+                    12,
+                    "RNN",
+                    10,
+                    10,
+                    n_epochs=1,
+                    torch_metrics=torch_metrics,
+                    metrics_params=metrics_params,
+                )
+                model.fit(self.series)