Make logging_metrics a kwarg

Author: Jan Beitner

pytorch_forecasting/models/base_model.py

@@ -13,6 +13,7 @@
 from pytorch_lightning.metrics.metric import TensorMetric
 from pytorch_lightning.utilities.parsing import get_init_args
 import torch
+import torch.nn as nn
 from torch.nn.utils import rnn
 from torch.optim.lr_scheduler import LambdaLR, OneCycleLR, ReduceLROnPlateau
 from torch.utils.data import DataLoader
@@ -57,7 +58,7 @@ def __init__(
         learning_rate: Union[float, List[float]] = 1e-3,
         log_gradient_flow: bool = False,
         loss: TensorMetric = SMAPE(),
-        logging_metrics: List[TensorMetric] = [],
+        logging_metrics: nn.ModuleList = nn.ModuleList([]),
         reduce_on_plateau_patience: int = 1000,
         weight_decay: float = 0.0,
         monotone_constaints: Dict[str, int] = {},
@@ -76,7 +77,8 @@ def __init__(
             log_gradient_flow (bool): If to log gradient flow, this takes time and should be only done to diagnose
                 training failures. Defaults to False.
             loss (TensorMetric, optional): metric to optimize. Defaults to SMAPE().
-            logging_metrics: (List[TensorMetric], optional): list of metrics to log.
+            logging_metrics (nn.ModuleList[MultiHorizonMetric]): list of metrics that are logged during training.
+                Defaults to [].
             reduce_on_plateau_patience (int): patience after which learning rate is reduced by a factor of 10. Defaults
                 to 1000
             weight_decay (float): weight decay. Defaults to 0.0.
@@ -102,7 +104,7 @@ def __init__(
         if not hasattr(self, "loss"):
             self.loss = loss
         if not hasattr(self, "logging_metrics"):
-            self.logging_metrics = logging_metrics
+            self.logging_metrics = nn.ModuleList([l for l in logging_metrics])
         if not hasattr(self, "output_transformer"):
             self.output_transformer = output_transformer
 

pytorch_forecasting/models/nbeats/__init__.py

@@ -33,6 +33,7 @@ def __init__(
         loss=SMAPE(),
         reduce_on_plateau_patience: int = 1000,
         backcast_loss_ratio: float = 0.0,
+        logging_metrics: nn.ModuleList = nn.ModuleList([SMAPE(), MAE(), RMSE(), MAPE(), MASE()]),
         **kwargs,
     ):
         """
@@ -76,11 +77,12 @@ def __init__(
             log_gradient_flow: if to log gradient flow, this takes time and should be only done to diagnose training
                 failures
             reduce_on_plateau_patience (int): patience after which learning rate is reduced by a factor of 10
+            logging_metrics (nn.ModuleList[MultiHorizonMetric]): list of metrics that are logged during training.
+                Defaults to nn.ModuleList([SMAPE(), MAE(), RMSE(), MAPE(), MASE()])
+            **kwargs: additional arguments to :py:class:`~BaseModel`.
         """
         self.save_hyperparameters()
-        self.logging_metrics = [SMAPE(), MAE(), RMSE(), MAPE(), MASE()]
-        super().__init__(**kwargs)
-        self.loss = loss
+        super().__init__(loss=loss, logging_metrics=logging_metrics, **kwargs)
 
         # setup stacks
         self.net_blocks = nn.ModuleList()

pytorch_forecasting/models/temporal_fusion_transformer/__init__.py

@@ -5,6 +5,7 @@
 
 from matplotlib import pyplot as plt
 import numpy as np
+from pytorch_lightning.metrics.metric import TensorMetric
 import torch
 from torch import nn
 from torch.nn.utils import rnn
@@ -55,7 +56,8 @@ def __init__(
         reduce_on_plateau_patience: int = 1000,
         monotone_constaints: Dict[str, int] = {},
         share_single_variable_networks: bool = False,
-        output_transformer: Callable = None,
+        logging_metrics: nn.ModuleList = nn.ModuleList([SMAPE(), MAE(), RMSE(), MAPE()]),
+        **kwargs,
     ):
         """
         Temporal Fusion Transformer for forecasting timeseries - use its :py:meth:`~from_dataset` method if possible.
@@ -122,14 +124,14 @@ def __init__(
                 This constraint significantly slows down training. Defaults to {}.
             share_single_variable_networks (bool): if to share the single variable networks between the encoder and
                 decoder. Defaults to False.
+            logging_metrics (nn.ModuleList[MultiHorizonMetric]): list of metrics that are logged during training.
+                Defaults to nn.ModuleList([SMAPE(), MAE(), RMSE(), MAPE()]).
+            **kwargs: additional arguments to :py:class:`~BaseModel`.
         """
         self.save_hyperparameters()
-        super().__init__()
         # store loss function separately as it is a module
         assert isinstance(loss, MultiHorizonMetric), "Loss has to of class `MultiHorizonMetric`"
-        self.loss = loss
-        self.output_transformer = output_transformer
-        self.logging_metrics = [SMAPE(), MAE(), RMSE(), MAPE(), MASE()]
+        super().__init__(loss=loss, logging_metrics=logging_metrics, **kwargs)
 
         # processing inputs
         # embeddings
@@ -201,7 +203,6 @@ def __init__(
         )
 
         # create single variable grns that are shared across decoder and encoder
-
         if self.hparams.share_single_variable_networks:
             self.shared_single_variable_grns = nn.ModuleDict()
             for name, input_size in encoder_input_sizes.items():