From 0d5c7220ff229700df03c2ff36e87a78646d0031 Mon Sep 17 00:00:00 2001
From: Cristof <45030708+cristof-r@users.noreply.github.com>
Date: Mon, 8 Apr 2024 18:07:26 +0200
Subject: [PATCH] Implement TSMixer Model (#2293)
---
.github/workflows/merge.yml | 2 +-
.gitignore | 1 +
CHANGELOG.md | 7 +-
README.md | 1 +
darts/models/__init__.py | 1 +
darts/models/forecasting/__init__.py | 1 +
darts/models/forecasting/tsmixer_model.py | 846 ++++++++++++++
.../test_global_forecasting_models.py | 11 +-
.../forecasting/test_historical_forecasts.py | 12 +
.../forecasting/test_probabilistic_models.py | 19 +
.../test_torch_forecasting_model.py | 3 +
.../tests/models/forecasting/test_tsmixer.py | 371 ++++++
docs/source/examples.rst | 10 +
docs/userguide/covariates.md | 1 +
docs/userguide/torch_forecasting_models.md | 1 +
examples/21-TSMixer-examples.ipynb | 1025 +++++++++++++++++
16 files changed, 2307 insertions(+), 5 deletions(-)
create mode 100644 darts/models/forecasting/tsmixer_model.py
create mode 100644 darts/tests/models/forecasting/test_tsmixer.py
create mode 100644 examples/21-TSMixer-examples.ipynb
diff --git a/.github/workflows/merge.yml b/.github/workflows/merge.yml
index d8b5ae732f..87bc82f63b 100644
--- a/.github/workflows/merge.yml
+++ b/.github/workflows/merge.yml
@@ -87,7 +87,7 @@ jobs:
runs-on: ubuntu-latest
strategy:
matrix:
- example-name: [00-quickstart.ipynb, 01-multi-time-series-and-covariates.ipynb, 02-data-processing.ipynb, 03-FFT-examples.ipynb, 04-RNN-examples.ipynb, 05-TCN-examples.ipynb, 06-Transformer-examples.ipynb, 07-NBEATS-examples.ipynb, 08-DeepAR-examples.ipynb, 09-DeepTCN-examples.ipynb, 10-Kalman-filter-examples.ipynb, 11-GP-filter-examples.ipynb, 12-Dynamic-Time-Warping-example.ipynb, 13-TFT-examples.ipynb, 15-static-covariates.ipynb, 16-hierarchical-reconciliation.ipynb, 18-TiDE-examples.ipynb, 19-EnsembleModel-examples.ipynb, 20-RegressionModel-examples.ipynb]
+ example-name: [00-quickstart.ipynb, 01-multi-time-series-and-covariates.ipynb, 02-data-processing.ipynb, 03-FFT-examples.ipynb, 04-RNN-examples.ipynb, 05-TCN-examples.ipynb, 06-Transformer-examples.ipynb, 07-NBEATS-examples.ipynb, 08-DeepAR-examples.ipynb, 09-DeepTCN-examples.ipynb, 10-Kalman-filter-examples.ipynb, 11-GP-filter-examples.ipynb, 12-Dynamic-Time-Warping-example.ipynb, 13-TFT-examples.ipynb, 15-static-covariates.ipynb, 16-hierarchical-reconciliation.ipynb, 18-TiDE-examples.ipynb, 19-EnsembleModel-examples.ipynb, 20-RegressionModel-examples.ipynb, 21-TSMixer-examples.ipynb]
steps:
- name: "1. Clone repository"
uses: actions/checkout@v2
diff --git a/.gitignore b/.gitignore
index 453913f0b7..1e3939db7f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -16,6 +16,7 @@ runs/
htmlcov
coverage.xml
.darts
+darts_logs/
docs_env
.DS_Store
.gradle
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b2e6da8fd9..258086b94e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,14 +9,15 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
### For users of the library:
**Improved**
-- πππ Improvements to metrics, historical forecasts, backtest, and residuals through major refactor. The refactor includes optimization of multiple process and improvemenets to consistency, reliability, and the documentation. Some of these necessary changes come at the cost of breaking changes. [#2284](https://github.com/unit8co/darts/pull/2284) by [Dennis Bader](https://github.com/dennisbader).
+- ππ New forecasting model: `TSMixerModel` as proposed in [this paper](https://arxiv.org/abs/2303.06053). An MLP based model that combines temporal, static and cross-sectional feature information using stacked mixing layers. [#1807](https://https://github.com/unit8co/darts/pull/001), by [Dennis Bader](https://github.com/dennisbader) and [Cristof Rojas](https://github.com/cristof-r).
+- ππ Improvements to metrics, historical forecasts, backtest, and residuals through major refactor. The refactor includes optimization of multiple process and improvemenets to consistency, reliability, and the documentation. Some of these necessary changes come at the cost of breaking changes. [#2284](https://github.com/unit8co/darts/pull/2284) by [Dennis Bader](https://github.com/dennisbader).
- Metrics:
- - Optimized all metrics, which now run >20 times faster than before for univariate series, and >>20 times for multivariate series. This boosts direct metric computations as well as backtesting and residuals computation!
+ - Optimized all metrics, which now run **> n * 20 times faster** than before for series with `n` components/columns. This boosts direct metric computations as well as backtesting and residuals computation!
- Added new metrics:
- Time aggregated metric `merr()` (Mean Error)
- Time aggregated scaled metrics `rmsse()`, and `msse()`: The (Root) Mean Squared Scaled Error.
- "Per time step" metrics that return a metric score per time step: `err()` (Error), `ae()` (Absolute Error), `se()` (Squared Error), `sle()` (Squared Log Error), `ase()` (Absolute Scaled Error), `sse` (Squared Scaled Error), `ape()` (Absolute Percentage Error), `sape()` (symmetric Absolute Percentage Error), `arre()` (Absolute Ranged Relative Error), `ql` (Quantile Loss)
- - All scaled metrics now accept `insample` series that can be overlapping into `pred_series` (before that had to end exactly one step before `pred_series`). Darts will handle the correct time extraction for you.
+ - All scaled metrics now accept `insample` series that can be overlapping into `pred_series` (before they had to end exactly one step before `pred_series`). Darts will handle the correct time extraction for you.
- Improvements to the documentation:
- Added a summary list of all metrics to the [metrics documentation page](https://unit8co.github.io/darts/generated_api/darts.metrics.html)
- Standardized the documentation of each metric (added formula, improved return documentation, ...)
diff --git a/README.md b/README.md
index 1786c968a6..4d482214cc 100644
--- a/README.md
+++ b/README.md
@@ -255,6 +255,7 @@ on bringing more models and features.
| [DLinearModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.dlinear.html#darts.models.forecasting.dlinear.DLinearModel) | [DLinear paper](https://arxiv.org/pdf/2205.13504.pdf) | π© π© | π© π© π© | π© π© | π© |
| [NLinearModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.nlinear.html#darts.models.forecasting.nlinear.NLinearModel) | [NLinear paper](https://arxiv.org/pdf/2205.13504.pdf) | π© π© | π© π© π© | π© π© | π© |
| [TiDEModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.tide_model.html#darts.models.forecasting.tide_model.TiDEModel) | [TiDE paper](https://arxiv.org/pdf/2304.08424.pdf) | π© π© | π© π© π© | π© π© | π© |
+| [TSMixerModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.tsmixer_model.html#darts.models.forecasting.tsmixer_model.TSMixerModel) | [TSMixer paper](https://arxiv.org/pdf/2303.06053.pdf), [PyTorch Implementation](https://github.com/ditschuk/pytorch-tsmixer) | π© π© | π© π© π© | π© π© | π© |
| **Ensemble Models**
([GlobalForecastingModel](https://unit8co.github.io/darts/userguide/covariates.html#global-forecasting-models-gfms)): Model support is dependent on ensembled forecasting models and the ensemble model itself | | | | | |
| [NaiveEnsembleModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.baselines.html#darts.models.forecasting.baselines.NaiveEnsembleModel) | | π© π© | π© π© π© | π© π© | π© |
| [RegressionEnsembleModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.regression_ensemble_model.html#darts.models.forecasting.regression_ensemble_model.RegressionEnsembleModel) | | π© π© | π© π© π© | π© π© | π© |
diff --git a/darts/models/__init__.py b/darts/models/__init__.py
index edcca507ea..3409aaa2ab 100644
--- a/darts/models/__init__.py
+++ b/darts/models/__init__.py
@@ -51,6 +51,7 @@
from darts.models.forecasting.tft_model import TFTModel
from darts.models.forecasting.tide_model import TiDEModel
from darts.models.forecasting.transformer_model import TransformerModel
+ from darts.models.forecasting.tsmixer_model import TSMixerModel
except ModuleNotFoundError:
logger.warning(
"Support for Torch based models not available. "
diff --git a/darts/models/forecasting/__init__.py b/darts/models/forecasting/__init__.py
index 9fa591ca27..37a50aa4bc 100644
--- a/darts/models/forecasting/__init__.py
+++ b/darts/models/forecasting/__init__.py
@@ -46,6 +46,7 @@
- :class:`~darts.models.forecasting.dlinear.DLinearModel`
- :class:`~darts.models.forecasting.nlinear.NLinearModel`
- :class:`~darts.models.forecasting.tide_model.TiDEModel`
+ - :class:`~darts.models.forecasting.tsmixer_model.TSMixerModel`
Ensemble Models (`GlobalForecastingModel `_)
- :class:`~darts.models.forecasting.baselines.NaiveEnsembleModel`
- :class:`~darts.models.forecasting.regression_ensemble_model.RegressionEnsembleModel`
diff --git a/darts/models/forecasting/tsmixer_model.py b/darts/models/forecasting/tsmixer_model.py
new file mode 100644
index 0000000000..0e53080739
--- /dev/null
+++ b/darts/models/forecasting/tsmixer_model.py
@@ -0,0 +1,846 @@
+"""
+Time-Series Mixer (TSMixer)
+---------------------------
+"""
+
+# The inner layers (``nn.Modules``) and the ``TimeBatchNorm2d`` were provided by a PyTorch implementation
+# of TSMixer: https://github.com/ditschuk/pytorch-tsmixer
+#
+# The License of pytorch-tsmixer v0.2.0 from https://github.com/ditschuk/pytorch-tsmixer/blob/main/LICENSE,
+# accessed Thursday, March 21st, 2024:
+# 'The MIT License
+#
+# Copyright 2023 Konstantin Ditschuneit
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
+# associated documentation files (the βSoftwareβ), to deal in the Software without restriction,
+# including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so,
+# subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all copies or substantial
+# portions of the Software.
+# '
+
+from typing import Callable, Optional, Tuple, Union
+
+import torch
+from torch import nn
+
+from darts.logging import get_logger, raise_log
+from darts.models.components import layer_norm_variants
+from darts.models.forecasting.pl_forecasting_module import (
+ PLMixedCovariatesModule,
+ io_processor,
+)
+from darts.models.forecasting.torch_forecasting_model import MixedCovariatesTorchModel
+from darts.utils.torch import MonteCarloDropout
+
+MixedCovariatesTrainTensorType = Tuple[
+ torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
+]
+
+logger = get_logger(__name__)
+
+ACTIVATIONS = [
+ "ReLU",
+ "RReLU",
+ "PReLU",
+ "ELU",
+ "Softplus",
+ "Tanh",
+ "SELU",
+ "LeakyReLU",
+ "Sigmoid",
+ "GELU",
+]
+
+NORMS = [
+ "LayerNorm",
+ "LayerNormNoBias",
+ "TimeBatchNorm2d",
+]
+
+
+def _time_to_feature(x: torch.Tensor) -> torch.Tensor:
+ """Converts a time series Tensor to a feature Tensor."""
+ return x.permute(0, 2, 1)
+
+
+class TimeBatchNorm2d(nn.BatchNorm2d):
+ def __init__(self, *args, **kwargs):
+ """A batch normalization layer that normalizes over the last two dimensions of a Tensor."""
+ super().__init__(num_features=1)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ # `x` has shape (batch_size, time, features)
+ if x.ndim != 3:
+ raise_log(
+ ValueError(
+ f"Expected 3D input Tensor, but got {x.ndim}D Tensor" " instead."
+ ),
+ logger=logger,
+ )
+ # apply 2D batch norm over reshape input_data `(batch_size, 1, timepoints, features)`
+ output = super().forward(x.unsqueeze(1))
+ # reshape back to (batch_size, timepoints, features)
+ return output.squeeze(1)
+
+
+class _FeatureMixing(nn.Module):
+ def __init__(
+ self,
+ sequence_length: int,
+ input_dim: int,
+ output_dim: int,
+ ff_size: int,
+ activation: Callable[[torch.Tensor], torch.Tensor],
+ dropout: float,
+ normalize_before: bool,
+ norm_type: nn.Module,
+ ) -> None:
+ """A module for feature mixing with flexibility in normalization and activation based on the
+ `PyTorch implementation of TSMixer `_.
+
+ This module provides options for batch normalization before or after mixing
+ features, uses dropout for regularization, and allows for different activation
+ functions.
+
+ Parameters
+ ----------
+ sequence_length
+ The length of the input sequences.
+ input_dim
+ The number of input channels to the module.
+ output_dim
+ The number of output channels from the module.
+ ff_size
+ The dimension of the feed-forward network internal to the module.
+ activation
+ The activation function used within the feed-forward network.
+ dropout
+ The dropout probability used for regularization.
+ normalize_before
+ A boolean indicating whether to apply normalization before
+ the rest of the operations.
+ norm_type
+ The type of normalization to use.
+ """
+ super().__init__()
+
+ self.projection = (
+ nn.Linear(input_dim, output_dim)
+ if input_dim != output_dim
+ else nn.Identity()
+ )
+ self.norm_before = (
+ norm_type((sequence_length, input_dim))
+ if normalize_before
+ else nn.Identity()
+ )
+ self.fc1 = nn.Linear(input_dim, ff_size)
+ self.activation = activation
+ self.dropout1 = MonteCarloDropout(dropout)
+ self.fc2 = nn.Linear(ff_size, output_dim)
+ self.dropout2 = MonteCarloDropout(dropout)
+ self.norm_after = (
+ norm_type((sequence_length, output_dim))
+ if not normalize_before
+ else nn.Identity()
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x_proj = self.projection(x)
+ x = self.norm_before(x)
+ x = self.fc1(x)
+ x = self.activation(x)
+ x = self.dropout1(x)
+ x = self.fc2(x)
+ x = self.dropout2(x)
+ x = x_proj + x
+ x = self.norm_after(x)
+ return x
+
+
+class _TimeMixing(nn.Module):
+ def __init__(
+ self,
+ sequence_length: int,
+ input_dim: int,
+ activation: Callable,
+ dropout: float,
+ normalize_before: bool,
+ norm_type: nn.Module,
+ ) -> None:
+ """Applies a transformation over the time dimension of a sequence based on the
+ `PyTorch implementation of TSMixer `_.
+
+ This module applies a linear transformation followed by an activation function
+ and dropout over the sequence length of the input feature torch.Tensor after converting
+ feature maps to the time dimension and then back.
+
+ Parameters
+ ----------
+ sequence_length
+ The length of the sequences to be transformed.
+ input_dim
+ The number of input channels to the module.
+ activation
+ The activation function to be used after the linear
+ transformation.
+ dropout
+ The dropout probability to be used after the activation function.
+ normalize_before
+ Whether to apply normalization before or after feature mixing.
+ norm_type
+ The type of normalization to use.
+ """
+ super().__init__()
+ self.normalize_before = normalize_before
+ self.norm_before = (
+ norm_type((sequence_length, input_dim))
+ if normalize_before
+ else nn.Identity()
+ )
+ self.activation = activation
+ self.dropout = MonteCarloDropout(dropout)
+ self.fc1 = nn.Linear(sequence_length, sequence_length)
+ self.norm_after = (
+ norm_type((sequence_length, input_dim))
+ if not normalize_before
+ else nn.Identity()
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ # permute the feature dim with the time dim
+ x_temp = self.norm_before(x)
+ x_temp = _time_to_feature(x_temp)
+ x_temp = self.activation(self.fc1(x_temp))
+ x_temp = self.dropout(x_temp)
+ # permute back the time dim with the feature dim
+ x_temp = x + _time_to_feature(x_temp)
+ x_temp = self.norm_after(x_temp)
+ return x_temp
+
+
+class _ConditionalMixerLayer(nn.Module):
+ def __init__(
+ self,
+ sequence_length: int,
+ input_dim: int,
+ output_dim: int,
+ static_cov_dim: int,
+ ff_size: int,
+ activation: Callable,
+ dropout: float,
+ normalize_before: bool,
+ norm_type: nn.Module,
+ ) -> None:
+ """Conditional mix layer combining time and feature mixing with static context based on the
+ `PyTorch implementation of TSMixer `_.
+
+ This module combines time mixing and conditional feature mixing, where the latter
+ is influenced by static features. This allows the module to learn representations
+ that are influenced by both dynamic and static features.
+
+ Parameters
+ ----------
+ sequence_length
+ The length of the input sequences.
+ input_dim
+ The number of input channels of the dynamic features.
+ output_dim
+ The number of output channels after feature mixing.
+ static_cov_dim
+ The number of channels in the static feature input.
+ ff_size
+ The inner dimension of the feedforward network used in feature mixing.
+ activation
+ The activation function used in both mixing operations.
+ dropout
+ The dropout probability used in both mixing operations.
+ normalize_before
+ Whether to apply normalization before or after mixing.
+ norm_type
+ The type of normalization to use.
+ """
+ super().__init__()
+
+ mixing_input = input_dim
+ if static_cov_dim != 0:
+ self.feature_mixing_static = _FeatureMixing(
+ sequence_length=sequence_length,
+ input_dim=static_cov_dim,
+ output_dim=output_dim,
+ ff_size=ff_size,
+ activation=activation,
+ dropout=dropout,
+ normalize_before=normalize_before,
+ norm_type=norm_type,
+ )
+ mixing_input += output_dim
+ else:
+ self.feature_mixing_static = None
+
+ self.time_mixing = _TimeMixing(
+ sequence_length=sequence_length,
+ input_dim=mixing_input,
+ activation=activation,
+ dropout=dropout,
+ normalize_before=normalize_before,
+ norm_type=norm_type,
+ )
+ self.feature_mixing = _FeatureMixing(
+ sequence_length=sequence_length,
+ input_dim=mixing_input,
+ output_dim=output_dim,
+ ff_size=ff_size,
+ activation=activation,
+ dropout=dropout,
+ normalize_before=normalize_before,
+ norm_type=norm_type,
+ )
+
+ def forward(
+ self, x: torch.Tensor, x_static: Optional[torch.Tensor]
+ ) -> torch.Tensor:
+ if self.feature_mixing_static is not None:
+ x_static_mixed = self.feature_mixing_static(x_static)
+ x = torch.cat([x, x_static_mixed], dim=-1)
+ x = self.time_mixing(x)
+ x = self.feature_mixing(x)
+ return x
+
+
+class _TSMixerModule(PLMixedCovariatesModule):
+ def __init__(
+ self,
+ input_dim: int,
+ output_dim: int,
+ past_cov_dim: int,
+ future_cov_dim: int,
+ static_cov_dim: int,
+ nr_params: int,
+ hidden_size: int,
+ ff_size: int,
+ num_blocks: int,
+ activation: str,
+ dropout: float,
+ norm_type: Union[str, nn.Module],
+ normalize_before: bool,
+ **kwargs,
+ ) -> None:
+ """
+ Initializes the TSMixer module for use within a Darts forecasting model.
+
+ Parameters
+ ----------
+ input_dim
+ Number of input target features.
+ output_dim
+ Number of output target features.
+ past_cov_dim
+ Number of past covariate features.
+ future_cov_dim
+ Number of future covariate features.
+ static_cov_dim
+ Number of static covariate features (number of target features
+ (or 1 if global static covariates) * number of static covariate features).
+ nr_params
+ The number of parameters of the likelihood (or 1 if no likelihood is used).
+ hidden_size
+ Hidden state size of the TSMixer.
+ ff_size
+ Dimension of the feedforward network internal to the module.
+ num_blocks
+ Number of mixer blocks.
+ activation
+ Activation function to use.
+ dropout
+ Dropout rate for regularization.
+ norm_type
+ Type of normalization to use.
+ normalize_before
+ Whether to apply normalization before or after mixing.
+ """
+ super().__init__(**kwargs)
+ self.input_dim = input_dim
+ self.output_dim = output_dim
+ self.future_cov_dim = future_cov_dim
+ self.static_cov_dim = static_cov_dim
+ self.nr_params = nr_params
+
+ if activation not in ACTIVATIONS:
+ raise_log(
+ ValueError(
+ f"Invalid `activation={activation}`. Must be on of {ACTIVATIONS}."
+ ),
+ logger=logger,
+ )
+ activation = getattr(nn, activation)()
+
+ if isinstance(norm_type, str):
+ if norm_type not in NORMS:
+ raise_log(
+ ValueError(
+ f"Invalid `norm_type={norm_type}`. Must be on of {NORMS}."
+ ),
+ logger=logger,
+ )
+ if norm_type == "TimeBatchNorm2d":
+ norm_type = TimeBatchNorm2d
+ else:
+ norm_type = getattr(layer_norm_variants, norm_type)
+ else:
+ norm_type = norm_type
+
+ mixer_params = {
+ "ff_size": ff_size,
+ "activation": activation,
+ "dropout": dropout,
+ "norm_type": norm_type,
+ "normalize_before": normalize_before,
+ }
+
+ self.fc_hist = nn.Linear(self.input_chunk_length, self.output_chunk_length)
+ self.feature_mixing_hist = _FeatureMixing(
+ sequence_length=self.output_chunk_length,
+ input_dim=input_dim + past_cov_dim + future_cov_dim,
+ output_dim=hidden_size,
+ **mixer_params,
+ )
+ if future_cov_dim:
+ self.feature_mixing_future = _FeatureMixing(
+ sequence_length=self.output_chunk_length,
+ input_dim=future_cov_dim,
+ output_dim=hidden_size,
+ **mixer_params,
+ )
+ else:
+ self.feature_mixing_future = None
+ self.conditional_mixer = self._build_mixer(
+ prediction_length=self.output_chunk_length,
+ num_blocks=num_blocks,
+ hidden_size=hidden_size,
+ future_cov_dim=future_cov_dim,
+ static_cov_dim=static_cov_dim,
+ **mixer_params,
+ )
+ self.fc_out = nn.Linear(hidden_size, output_dim * nr_params)
+
+ @staticmethod
+ def _build_mixer(
+ prediction_length: int,
+ num_blocks: int,
+ hidden_size: int,
+ future_cov_dim: int,
+ static_cov_dim: int,
+ **kwargs,
+ ) -> nn.ModuleList:
+ """Build the mixer blocks for the model."""
+ # the first block takes `x` consisting of concatenated features with size `hidden_size`:
+ # - historic features
+ # - optional future features
+ input_dim_block = hidden_size * (1 + int(future_cov_dim > 0))
+
+ mixer_layers = nn.ModuleList()
+ for _ in range(num_blocks):
+ layer = _ConditionalMixerLayer(
+ input_dim=input_dim_block,
+ output_dim=hidden_size,
+ sequence_length=prediction_length,
+ static_cov_dim=static_cov_dim,
+ **kwargs,
+ )
+ mixer_layers.append(layer)
+ # after the first block, `x` consists of previous block output with size `hidden_size`
+ input_dim_block = hidden_size
+ return mixer_layers
+
+ @io_processor
+ def forward(
+ self,
+ x_in: Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]],
+ ) -> torch.Tensor:
+ # x_hist contains the historical time series data and the historical
+ """TSMixer model forward pass.
+
+ Parameters
+ ----------
+ x_in
+ comes as Tuple `(x_past, x_future, x_static)` where `x_past` is the input/past chunk and
+ `x_future` is the output/future chunk. Input dimensions are `(batch_size, time_steps,
+ components)`.
+
+ Returns
+ -------
+ torch.torch.Tensor
+ The output Tensorof shape `(batch_size, output_chunk_length, output_dim, nr_params)`.
+ """
+ # B: batch size
+ # L: input chunk length
+ # T: output chunk length
+ # C: target components
+ # P: past cov features
+ # F: future cov features
+ # S: static cov features
+ # H = C + P + F: historic features
+ # H_S: hidden Size
+ # N_P: likelihood parameters
+
+ # `x`: (B, L, H), `x_future`: (B, T, F), `x_static`: (B, C or 1, S)
+ x, x_future, x_static = x_in
+
+ # swap feature and time dimensions (B, L, H) -> (B, H, L)
+ x = _time_to_feature(x)
+ # linear transformations to horizon (B, H, L) -> (B, H, T)
+ x = self.fc_hist(x)
+ # (B, H, T) -> (B, T, H)
+ x = _time_to_feature(x)
+
+ # feature mixing for historical features (B, T, H) -> (B, T, H_S)
+ x = self.feature_mixing_hist(x)
+ if self.future_cov_dim:
+ # feature mixing for future features (B, T, F) -> (B, T, H_S)
+ x_future = self.feature_mixing_future(x_future)
+ # (B, T, H_S) + (B, T, H_S) -> (B, T, 2*H_S)
+ x = torch.cat([x, x_future], dim=-1)
+
+ if self.static_cov_dim:
+ # (B, C, S) -> (B, 1, C * S)
+ x_static = x_static.reshape(x_static.shape[0], 1, -1)
+ # repeat to match horizon (B, 1, C * S) -> (B, T, C * S)
+ x_static = x_static.repeat(1, self.output_chunk_length, 1)
+
+ for mixing_layer in self.conditional_mixer:
+ # conditional mixer layers with static covariates (B, T, 2 * H_S), (B, T, C * S) -> (B, T, H_S)
+ x = mixing_layer(x, x_static=x_static)
+
+ # linear transformation to generate the forecast (B, T, H_S) -> (B, T, C * N_P)
+ x = self.fc_out(x)
+ # (B, T, C * N_P) -> (B, T, C, N_P)
+ x = x.view(-1, self.output_chunk_length, self.output_dim, self.nr_params)
+ return x
+
+
+class TSMixerModel(MixedCovariatesTorchModel):
+ def __init__(
+ self,
+ input_chunk_length: int,
+ output_chunk_length: int,
+ output_chunk_shift: int = 0,
+ hidden_size: int = 64,
+ ff_size: int = 64,
+ num_blocks: int = 2,
+ activation: str = "ReLU",
+ dropout: float = 0.1,
+ norm_type: Union[str, nn.Module] = "LayerNorm",
+ normalize_before: bool = False,
+ use_static_covariates: bool = True,
+ **kwargs,
+ ) -> None:
+ """Time-Series Mixer (TSMixer): An All-MLP Architecture for Time Series.
+
+ This is an implementation of the TSMixer architecture, as outlined in [1]_. A major part of the architecture
+ was adopted from `this PyTorch implementation `_. Additional
+ changes were applied to increase model performance and efficiency.
+
+ TSMixer forecasts time series data by integrating historical time series data, future known inputs, and static
+ contextual information. It uses a combination of conditional feature mixing and mixer layers to process and
+ combine these different types of data for effective forecasting.
+
+ This model supports past covariates (known for `input_chunk_length` points before prediction time), future
+ covariates (known for `output_chunk_length` points after prediction time), static covariates, as well as
+ probabilistic forecasting.
+
+ Parameters
+ ----------
+ input_chunk_length
+ Number of time steps in the past to take as a model input (per chunk). Applies to the target
+ series, and past and/or future covariates (if the model supports it).
+ Also called: Encoder length
+ output_chunk_length
+ Number of time steps predicted at once (per chunk) by the internal model. Also, the number of future values
+ from future covariates to use as a model input (if the model supports future covariates). It is not the same
+ as forecast horizon `n` used in `predict()`, which is the desired number of prediction points generated
+ using either a one-shot- or autoregressive forecast. Setting `n <= output_chunk_length` prevents
+ auto-regression. This is useful when the covariates don't extend far enough into the future, or to prohibit
+ the model from using future values of past and / or future covariates for prediction (depending on the
+ model's covariate support).
+ Also called: Decoder length
+ output_chunk_shift
+ Optionally, the number of steps to shift the start of the output chunk into the future (relative to the
+ input chunk end). This will create a gap between the input and output. If the model supports
+ `future_covariates`, the future values are extracted from the shifted output chunk. Predictions will start
+ `output_chunk_shift` steps after the end of the target `series`. If `output_chunk_shift` is set, the model
+ cannot generate autoregressive predictions (`n > output_chunk_length`).
+ hidden_size
+ The hidden state size / size of the second feed-forward layer in the feature mixing MLP.
+ ff_size
+ The size of the first feed-forward layer in the feature mixing MLP.
+ num_blocks
+ The number of mixer blocks in the model. The number includes the first block and all subsequent blocks.
+ activation
+ The name of the activation function to use in the mixer layers. Default: `"ReLU"`. Must be one of
+ `"ReLU", "RReLU", "PReLU", "ELU", "Softplus", "Tanh", "SELU", "LeakyReLU", "Sigmoid", "GELU"`.
+ dropout
+ Fraction of neurons affected by dropout. This is compatible with Monte Carlo dropout at inference time
+ for model uncertainty estimation (enabled with ``mc_dropout=True`` at prediction time).
+ norm_type
+ The type of `LayerNorm` variant to use. Default: `"LayerNorm"`. If a string, must be one of
+ `"LayerNormNoBias", "LayerNorm", "TimeBatchNorm2d"`. Otherwise, must be a custom `nn.Module`.
+ normalize_before
+ Whether to apply layer normalization before or after mixer layer.
+ use_static_covariates
+ Whether the model should use static covariate information in case the input `series` passed to ``fit()``
+ contain static covariates. If ``True``, and static covariates are available at fitting time, will enforce
+ that all target `series` have the same static covariate dimensionality in ``fit()`` and ``predict()``.
+ **kwargs
+ Optional arguments to initialize the pytorch_lightning.Module, pytorch_lightning.Trainer, and
+ Darts' :class:`TorchForecastingModel`.
+
+ loss_fn
+ 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.
+ likelihood
+ The likelihood model to be used for probabilistic forecasts.
+ 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``.
+ optimizer_cls
+ The PyTorch optimizer class to be used. Default: ``torch.optim.Adam``.
+ optimizer_kwargs
+ Optionally, some keyword arguments for the PyTorch optimizer (e.g., ``{'lr': 1e-3}``
+ for specifying a learning rate). Otherwise, the default values of the selected ``optimizer_cls``
+ will be used. Default: ``None``.
+ lr_scheduler_cls
+ Optionally, the PyTorch learning rate scheduler class to be used. Specifying ``None`` corresponds
+ to using a constant learning rate. Default: ``None``.
+ lr_scheduler_kwargs
+ Optionally, some keyword arguments for the PyTorch learning rate scheduler. Default: ``None``.
+ use_reversible_instance_norm
+ Whether to use reversible instance normalization `RINorm` against distribution shift as shown in [3]_.
+ It is only applied to the features of the target series and not the covariates.
+ batch_size
+ Number of time series (input and output sequences) used in each training pass. Default: ``32``.
+ n_epochs
+ Number of epochs over which to train the model. Default: ``100``.
+ model_name
+ Name of the model. Used for creating checkpoints and saving torch.Tensorboard data. If not specified,
+ defaults to the following string ``"YYYY-mm-dd_HH_MM_SS_torch_model_run_PID"``, where the initial part
+ of the name is formatted with the local date and time, while PID is the processed ID (preventing models
+ spawned at the same time by different processes to share the same model_name). E.g.,
+ ``"2021-06-14_09_53_32_torch_model_run_44607"``.
+ work_dir
+ Path of the working directory, where to save checkpoints and torch.Tensorboard summaries.
+ Default: current working directory.
+ log_torch.Tensorboard
+ If set, use torch.Tensorboard to log the different parameters. The logs will be located in:
+ ``"{work_dir}/darts_logs/{model_name}/logs/"``. Default: ``False``.
+ nr_epochs_val_period
+ Number of epochs to wait before evaluating the validation loss (if a validation
+ ``TimeSeries`` is passed to the :func:`fit()` method). Default: ``1``.
+ force_reset
+ If set to ``True``, any previously-existing model with the same name will be reset (all checkpoints will
+ be discarded). Default: ``False``.
+ save_checkpoints
+ Whether to automatically save the untrained model and checkpoints from training.
+ To load the model from checkpoint, call :func:`MyModelClass.load_from_checkpoint()`, where
+ :class:`MyModelClass` is the :class:`TorchForecastingModel` class that was used (such as :class:`TFTModel`,
+ :class:`NBEATSModel`, etc.). If set to ``False``, the model can still be manually saved using
+ :func:`save()` and loaded using :func:`load()`. Default: ``False``.
+ add_encoders
+ A large number of past and future covariates can be automatically generated with `add_encoders`.
+ This can be done by adding multiple pre-defined index encoders and/or custom user-made functions that
+ will be used as index encoders. Additionally, a transformer such as Darts' :class:`Scaler` can be added to
+ transform the generated covariates. This happens all under one hood and only needs to be specified at
+ model creation.
+ Read :meth:`SequentialEncoder ` to find out more about
+ ``add_encoders``. Default: ``None``. An example showing some of ``add_encoders`` features:
+
+ .. highlight:: python
+ .. code-block:: python
+
+ def encode_year(idx):
+ return (idx.year - 1950) / 50
+
+ add_encoders={
+ 'cyclic': {'future': ['month']},
+ 'datetime_attribute': {'future': ['hour', 'dayofweek']},
+ 'position': {'past': ['relative'], 'future': ['relative']},
+ 'custom': {'past': [encode_year]},
+ 'transformer': Scaler(),
+ 'tz': 'CET'
+ }
+ ..
+ random_state
+ Control the randomness of the weight's initialization. Check this
+ `link `_ for more details.
+ Default: ``None``.
+ pl_trainer_kwargs
+ By default :class:`TorchForecastingModel` creates a PyTorch Lightning Trainer with several useful presets
+ that performs the training, validation and prediction processes. These presets include automatic
+ checkpointing, torch.Tensorboard logging, setting the torch device and more.
+ With ``pl_trainer_kwargs`` you can add additional kwargs to instantiate the PyTorch Lightning trainer
+ object. Check the `PL Trainer documentation
+ `_ for more information about the
+ supported kwargs. Default: ``None``.
+ Running on GPU(s) is also possible using ``pl_trainer_kwargs`` by specifying keys ``"accelerator",
+ "devices", and "auto_select_gpus"``. Some examples for setting the devices inside the ``pl_trainer_kwargs``
+ dict:
+
+ - ``{"accelerator": "cpu"}`` for CPU,
+ - ``{"accelerator": "gpu", "devices": [i]}`` to use only GPU ``i`` (``i`` must be an integer),
+ - ``{"accelerator": "gpu", "devices": -1, "auto_select_gpus": True}`` to use all available GPUS.
+
+ For more info, see here:
+ https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html#trainer-flags , and
+ https://pytorch-lightning.readthedocs.io/en/stable/accelerators/gpu_basic.html#train-on-multiple-gpus
+
+ With parameter ``"callbacks"`` you can add custom or PyTorch-Lightning built-in callbacks to Darts'
+ :class:`TorchForecastingModel`. Below is an example for adding EarlyStopping to the training process.
+ The model will stop training early if the validation loss `val_loss` does not improve beyond
+ specifications. For more information on callbacks, visit:
+ `PyTorch Lightning Callbacks
+ `_
+
+ .. highlight:: python
+ .. code-block:: python
+
+ from pytorch_lightning.callbacks.early_stopping import EarlyStopping
+
+ # stop training when validation loss does not decrease more than 0.05 (`min_delta`) over
+ # a period of 5 epochs (`patience`)
+ my_stopper = EarlyStopping(
+ monitor="val_loss",
+ patience=5,
+ min_delta=0.05,
+ mode='min',
+ )
+
+ pl_trainer_kwargs={"callbacks": [my_stopper]}
+ ..
+
+ Note that you can also use a custom PyTorch Lightning Trainer for training and prediction with optional
+ parameter ``trainer`` in :func:`fit()` and :func:`predict()`.
+ show_warnings
+ whether to show warnings raised from PyTorch Lightning. Useful to detect potential issues of
+ your forecasting use case. Default: ``False``.
+
+ References
+ ----------
+ .. [1] https://arxiv.org/abs/2303.06053
+
+ Examples
+ --------
+ >>> from darts.datasets import WeatherDataset
+ >>> from darts.models import TSMixerModel
+ >>> series = WeatherDataset().load()
+ >>> # predicting temperatures
+ >>> target = series['T (degC)'][:100]
+ >>> # optionally, use past observed rainfall (pretending to be unknown beyond index 100)
+ >>> past_cov = series['rain (mm)'][:100]
+ >>> # optionally, use future atmospheric pressure (pretending this component is a forecast)
+ >>> future_cov = series['p (mbar)'][:106]
+ >>> model = TSMixerModel(
+ >>> input_chunk_length=6,
+ >>> output_chunk_length=6,
+ >>> use_reversible_instance_norm=True,
+ >>> n_epochs=20
+ >>> )
+ >>> model.fit(target, past_covariates=past_cov, future_covariates=future_cov)
+ >>> pred = model.predict(6)
+ >>> pred.values()
+ array([[3.92519848],
+ [4.05650312],
+ [4.21781987],
+ [4.29394973],
+ [4.4122863 ],
+ [4.42762751]])
+ """
+ model_kwargs = {key: val for key, val in self.model_params.items()}
+ super().__init__(**self._extract_torch_model_params(**model_kwargs))
+
+ # extract pytorch lightning module kwargs
+ self.pl_module_params = self._extract_pl_module_params(**model_kwargs)
+
+ # Model specific parameters
+ self.ff_size = ff_size
+ self.dropout = dropout
+ self.num_blocks = num_blocks
+ self.activation = activation
+ self.normalize_before = normalize_before
+ self.norm_type = norm_type
+ self.hidden_size = hidden_size
+ self._considers_static_covariates = use_static_covariates
+
+ def _create_model(self, train_sample: MixedCovariatesTrainTensorType) -> nn.Module:
+ """
+ Parameters
+ ----------
+ train_sample
+ contains the following torch.Tensors: `(past_target, past_covariates, historic_future_covariates,
+ future_covariates, static_covariates, future_target)`:
+
+ - past/historic torch.Tensors have shape (input_chunk_length, n_variables)
+ - future torch.Tensors have shape (output_chunk_length, n_variables)
+ - static covariates have shape (component, static variable)
+ """
+ (
+ past_target,
+ past_covariates,
+ historic_future_covariates,
+ future_covariates,
+ static_covariates,
+ future_target,
+ ) = train_sample
+
+ input_dim = past_target.shape[1]
+ output_dim = future_target.shape[1]
+
+ static_cov_dim = (
+ static_covariates.shape[0] * static_covariates.shape[1]
+ if static_covariates is not None
+ else 0
+ )
+ future_cov_dim = (
+ future_covariates.shape[1] if future_covariates is not None else 0
+ )
+ past_cov_dim = past_covariates.shape[1] if past_covariates is not None else 0
+ nr_params = 1 if self.likelihood is None else self.likelihood.num_parameters
+
+ return _TSMixerModule(
+ input_dim=input_dim,
+ output_dim=output_dim,
+ future_cov_dim=future_cov_dim,
+ past_cov_dim=past_cov_dim,
+ static_cov_dim=static_cov_dim,
+ nr_params=nr_params,
+ hidden_size=self.hidden_size,
+ ff_size=self.ff_size,
+ num_blocks=self.num_blocks,
+ activation=self.activation,
+ dropout=self.dropout,
+ norm_type=self.norm_type,
+ normalize_before=self.normalize_before,
+ **self.pl_module_params,
+ )
+
+ @property
+ def supports_multivariate(self) -> bool:
+ return True
+
+ @property
+ def supports_static_covariates(self) -> bool:
+ return True
+
+ @property
+ def supports_future_covariates(self) -> bool:
+ return True
+
+ @property
+ def supports_past_covariates(self) -> bool:
+ return True
diff --git a/darts/tests/models/forecasting/test_global_forecasting_models.py b/darts/tests/models/forecasting/test_global_forecasting_models.py
index b8b020f342..dd3e6faf8d 100644
--- a/darts/tests/models/forecasting/test_global_forecasting_models.py
+++ b/darts/tests/models/forecasting/test_global_forecasting_models.py
@@ -33,6 +33,7 @@
TFTModel,
TiDEModel,
TransformerModel,
+ TSMixerModel,
)
from darts.models.forecasting.torch_forecasting_model import (
DualCovariatesTorchModel,
@@ -155,6 +156,14 @@
},
40.0,
),
+ (
+ TSMixerModel,
+ {
+ "n_epochs": 10,
+ "pl_trainer_kwargs": tfm_kwargs["pl_trainer_kwargs"],
+ },
+ 60.0,
+ ),
(
GlobalNaiveAggregate,
{
@@ -527,7 +536,7 @@ def test_future_covariates(self):
@pytest.mark.parametrize(
"model_cls,ts",
product(
- [TFTModel, DLinearModel, NLinearModel, TiDEModel],
+ [TFTModel, DLinearModel, NLinearModel, TiDEModel, TSMixerModel],
[ts_w_static_cov, ts_shared_static_cov, ts_comps_static_cov],
),
)
diff --git a/darts/tests/models/forecasting/test_historical_forecasts.py b/darts/tests/models/forecasting/test_historical_forecasts.py
index 236933b714..e92eedffdc 100644
--- a/darts/tests/models/forecasting/test_historical_forecasts.py
+++ b/darts/tests/models/forecasting/test_historical_forecasts.py
@@ -38,6 +38,7 @@
TFTModel,
TiDEModel,
TransformerModel,
+ TSMixerModel,
)
from darts.utils.likelihood_models import GaussianLikelihood, QuantileRegression
@@ -235,6 +236,17 @@
(IN_LEN, OUT_LEN),
"MixedCovariates",
),
+ (
+ TSMixerModel,
+ {
+ "input_chunk_length": IN_LEN,
+ "output_chunk_length": OUT_LEN,
+ "n_epochs": NB_EPOCH,
+ **tfm_kwargs,
+ },
+ (IN_LEN, OUT_LEN),
+ "MixedCovariates",
+ ),
(
GlobalNaiveAggregate,
{
diff --git a/darts/tests/models/forecasting/test_probabilistic_models.py b/darts/tests/models/forecasting/test_probabilistic_models.py
index 7b728efcbb..169f9b6a1e 100644
--- a/darts/tests/models/forecasting/test_probabilistic_models.py
+++ b/darts/tests/models/forecasting/test_probabilistic_models.py
@@ -35,6 +35,7 @@
TFTModel,
TiDEModel,
TransformerModel,
+ TSMixerModel,
)
from darts.models.forecasting.torch_forecasting_model import TorchForecastingModel
from darts.utils.likelihood_models import (
@@ -194,6 +195,24 @@
0.06,
0.1,
),
+ (
+ TSMixerModel,
+ {
+ "input_chunk_length": 10,
+ "output_chunk_length": 5,
+ "n_epochs": 100,
+ "random_state": 0,
+ "num_blocks": 1,
+ "hidden_size": 32,
+ "dropout": 0.2,
+ "ff_size": 32,
+ "batch_size": 8,
+ "likelihood": GaussianLikelihood(),
+ **tfm_kwargs,
+ },
+ 0.06,
+ 0.1,
+ ),
]
diff --git a/darts/tests/models/forecasting/test_torch_forecasting_model.py b/darts/tests/models/forecasting/test_torch_forecasting_model.py
index 73ec9bb19b..0e04f821b8 100644
--- a/darts/tests/models/forecasting/test_torch_forecasting_model.py
+++ b/darts/tests/models/forecasting/test_torch_forecasting_model.py
@@ -41,6 +41,7 @@
TFTModel,
TiDEModel,
TransformerModel,
+ TSMixerModel,
)
from darts.models.components.layer_norm_variants import RINorm
from darts.utils.likelihood_models import (
@@ -66,6 +67,7 @@
(TFTModel, {"add_relative_index": 2, **kwargs}),
(TiDEModel, kwargs),
(TransformerModel, kwargs),
+ (TSMixerModel, kwargs),
(GlobalNaiveSeasonal, kwargs),
(GlobalNaiveAggregate, kwargs),
(GlobalNaiveDrift, kwargs),
@@ -1505,6 +1507,7 @@ def test_rin(self, model_config):
(NHiTSModel, {}),
(TransformerModel, {}),
(TCNModel, {}),
+ (TSMixerModel, {}),
(BlockRNNModel, {}),
(GlobalNaiveSeasonal, {}),
(GlobalNaiveAggregate, {}),
diff --git a/darts/tests/models/forecasting/test_tsmixer.py b/darts/tests/models/forecasting/test_tsmixer.py
new file mode 100644
index 0000000000..6ae3abe39e
--- /dev/null
+++ b/darts/tests/models/forecasting/test_tsmixer.py
@@ -0,0 +1,371 @@
+from darts.logging import get_logger
+
+logger = get_logger(__name__)
+
+try:
+ import numpy as np
+ import pandas as pd
+ import pytest
+ import torch
+ from torch import nn
+
+ from darts import concatenate
+ from darts.models.forecasting.tsmixer_model import TimeBatchNorm2d, TSMixerModel
+ from darts.tests.conftest import tfm_kwargs
+ from darts.utils import timeseries_generation as tg
+ from darts.utils.likelihood_models import GaussianLikelihood
+
+ TORCH_AVAILABLE = True
+
+except ImportError:
+ logger.warning("Torch not available. TSMixerModel tests will be skipped.")
+ TORCH_AVAILABLE = False
+
+
+@pytest.mark.skipif(
+ TORCH_AVAILABLE is False,
+ reason="Torch not available. TSMixerModel tests will be skipped.",
+)
+class TestTSMixerModel:
+ np.random.seed(42)
+ torch.manual_seed(42)
+
+ def test_creation(self):
+ model = TSMixerModel(
+ input_chunk_length=1,
+ output_chunk_length=1,
+ likelihood=GaussianLikelihood(),
+ )
+
+ assert model.input_chunk_length == 1
+
+ def test_fit(self):
+ large_ts = tg.constant_timeseries(length=10, value=1.0)
+ small_ts = tg.constant_timeseries(length=10, value=0.1)
+
+ model = TSMixerModel(
+ input_chunk_length=1,
+ output_chunk_length=1,
+ n_epochs=10,
+ random_state=42,
+ **tfm_kwargs,
+ )
+
+ model.fit(large_ts)
+ pred = model.predict(n=2).values()[0]
+
+ # Test whether model trained on one series is better
+ # than one trained on another
+ model2 = TSMixerModel(
+ input_chunk_length=1,
+ output_chunk_length=1,
+ n_epochs=10,
+ random_state=42,
+ **tfm_kwargs,
+ )
+
+ model2.fit(small_ts)
+ pred2 = model2.predict(n=2).values()[0]
+ assert abs(pred2 - 0.1) < abs(pred - 0.1)
+
+ # test short predict
+ pred3 = model2.predict(n=1)
+ assert len(pred3) == 1
+
+ def test_likelihood_fit(self):
+ ts = tg.constant_timeseries(length=3)
+
+ model = TSMixerModel(
+ input_chunk_length=1,
+ output_chunk_length=1,
+ n_epochs=1,
+ random_state=42,
+ likelihood=GaussianLikelihood(),
+ **tfm_kwargs,
+ )
+ model.fit(ts)
+ # sampled from distribution
+ pred = model.predict(n=1, num_samples=20)
+ assert pred.n_samples == 20
+
+ # direct distribution parameter prediction
+ pred = model.predict(n=1, num_samples=1, predict_likelihood_parameters=True)
+ assert pred.n_components == 2
+ assert pred.n_samples == 1
+
+ model = TSMixerModel(
+ input_chunk_length=1,
+ output_chunk_length=1,
+ n_epochs=1,
+ random_state=42,
+ **tfm_kwargs,
+ )
+ model.fit(ts)
+ # mc dropout
+ pred = model.predict(n=1, mc_dropout=True, num_samples=10)
+ assert pred.n_samples == 10
+
+ def test_logtensorboard(self, tmpdir_module):
+ ts = tg.constant_timeseries(length=4)
+
+ # Test basic fit and predict
+ model = TSMixerModel(
+ input_chunk_length=1,
+ output_chunk_length=1,
+ n_epochs=1,
+ log_tensorboard=True,
+ batch_size=2,
+ work_dir=tmpdir_module,
+ pl_trainer_kwargs={
+ "log_every_n_steps": 1,
+ **tfm_kwargs["pl_trainer_kwargs"],
+ },
+ )
+ model.fit(ts)
+ _ = model.predict(n=2)
+
+ def test_static_covariates_support(self):
+ target_multi = concatenate(
+ [tg.sine_timeseries(length=10, freq="h")] * 2, axis=1
+ )
+
+ target_multi = target_multi.with_static_covariates(
+ pd.DataFrame(
+ [[0.0, 1.0, 0, 2], [2.0, 3.0, 1, 3]],
+ columns=["st1", "st2", "cat1", "cat2"],
+ )
+ )
+
+ # should work with cyclic encoding for time index
+ model = TSMixerModel(
+ input_chunk_length=3,
+ output_chunk_length=4,
+ add_encoders={"cyclic": {"future": "hour", "past": "hour"}},
+ pl_trainer_kwargs={
+ "fast_dev_run": True,
+ **tfm_kwargs["pl_trainer_kwargs"],
+ },
+ )
+ model.fit(target_multi, verbose=False)
+
+ assert model.model.static_cov_dim == np.prod(
+ target_multi.static_covariates.values.shape
+ )
+
+ # raise an error when trained with static covariates of wrong dimensionality
+ target_multi = target_multi.with_static_covariates(
+ pd.concat([target_multi.static_covariates] * 2, axis=1)
+ )
+ with pytest.raises(ValueError):
+ model.predict(n=1, series=target_multi, verbose=False)
+
+ # raise an error when trained with static covariates and trying to predict without
+ with pytest.raises(ValueError):
+ model.predict(
+ n=1, series=target_multi.with_static_covariates(None), verbose=False
+ )
+
+ # with `use_static_covariates=False`, we can predict without static covs
+ model = TSMixerModel(
+ input_chunk_length=3,
+ output_chunk_length=4,
+ use_static_covariates=False,
+ n_epochs=1,
+ **tfm_kwargs,
+ )
+ model.fit(target_multi)
+ preds = model.predict(n=2, series=target_multi.with_static_covariates(None))
+ assert preds.static_covariates is None
+
+ model = TSMixerModel(
+ input_chunk_length=3,
+ output_chunk_length=4,
+ use_static_covariates=False,
+ n_epochs=1,
+ **tfm_kwargs,
+ )
+ model.fit(target_multi.with_static_covariates(None))
+ preds = model.predict(n=2, series=target_multi)
+ assert preds.static_covariates.equals(target_multi.static_covariates)
+
+ @pytest.mark.parametrize("enable_rin", [True, False])
+ def test_future_covariate_handling(self, enable_rin):
+ ts_time_index = tg.sine_timeseries(length=2, freq="h")
+
+ model = TSMixerModel(
+ input_chunk_length=1,
+ output_chunk_length=1,
+ add_encoders={"cyclic": {"future": "hour"}},
+ use_reversible_instance_norm=enable_rin,
+ **tfm_kwargs,
+ )
+ model.fit(ts_time_index, verbose=False, epochs=1)
+
+ def test_past_covariate_handling(self):
+ ts_time_index = tg.sine_timeseries(length=2, freq="h")
+
+ model = TSMixerModel(
+ input_chunk_length=1,
+ output_chunk_length=1,
+ add_encoders={"cyclic": {"past": "hour"}},
+ **tfm_kwargs,
+ )
+ model.fit(ts_time_index, verbose=False, epochs=1)
+
+ def test_future_and_past_covariate_handling(self):
+ ts_time_index = tg.sine_timeseries(length=2, freq="h")
+
+ model = TSMixerModel(
+ input_chunk_length=1,
+ output_chunk_length=1,
+ add_encoders={"cyclic": {"future": "hour", "past": "hour"}},
+ **tfm_kwargs,
+ )
+ model.fit(ts_time_index, verbose=False, epochs=1)
+
+ def test_future_past_and_static_covariate_as_timeseries_handling(self):
+ ts_time_index = tg.sine_timeseries(length=2, freq="h")
+ ts_time_index = ts_time_index.with_static_covariates(
+ pd.DataFrame(
+ [
+ [
+ 0.0,
+ ]
+ ],
+ columns=["st1"],
+ )
+ )
+ for enable_rin in [True, False]:
+ # test with past_covariates timeseries
+ model = TSMixerModel(
+ input_chunk_length=1,
+ output_chunk_length=1,
+ add_encoders={"cyclic": {"future": "hour"}},
+ use_reversible_instance_norm=enable_rin,
+ **tfm_kwargs,
+ )
+ model.fit(
+ ts_time_index,
+ past_covariates=ts_time_index,
+ verbose=False,
+ epochs=1,
+ )
+
+ # test with past_covariates and future_covariates timeseries
+ model = TSMixerModel(
+ input_chunk_length=1,
+ output_chunk_length=1,
+ add_encoders={"cyclic": {"future": "hour", "past": "hour"}},
+ use_reversible_instance_norm=enable_rin,
+ **tfm_kwargs,
+ )
+ model.fit(
+ ts_time_index,
+ past_covariates=ts_time_index,
+ future_covariates=ts_time_index,
+ verbose=False,
+ epochs=1,
+ )
+
+ @pytest.mark.parametrize(
+ "norm_type, expect_exception",
+ [
+ ("LayerNorm", False),
+ ("LayerNormNoBias", False),
+ (nn.LayerNorm, False),
+ ("TimeBatchNorm2d", False),
+ ("invalid", True),
+ ],
+ )
+ def test_layer_norms_with_parametrization(self, norm_type, expect_exception):
+ series = tg.sine_timeseries(length=3)
+ base_model = TSMixerModel
+
+ if expect_exception:
+ with pytest.raises(ValueError):
+ model = base_model(
+ input_chunk_length=1,
+ output_chunk_length=1,
+ norm_type=norm_type,
+ **tfm_kwargs,
+ )
+ model.fit(series, epochs=1)
+ else:
+ model = base_model(
+ input_chunk_length=1,
+ output_chunk_length=1,
+ norm_type=norm_type,
+ **tfm_kwargs,
+ )
+ model.fit(series, epochs=1)
+
+ @pytest.mark.parametrize(
+ "activation, expect_error",
+ [
+ ("ReLU", False),
+ ("RReLU", False),
+ ("PReLU", False),
+ ("ELU", False),
+ ("Softplus", False),
+ ("Tanh", False),
+ ("SELU", False),
+ ("LeakyReLU", False),
+ ("Sigmoid", False),
+ ("invalid", True),
+ ],
+ )
+ def test_activation_functions(self, activation, expect_error):
+ series = tg.sine_timeseries(length=3)
+ base_model = TSMixerModel
+
+ if expect_error:
+ with pytest.raises(ValueError):
+ model = base_model(
+ input_chunk_length=1,
+ output_chunk_length=1,
+ activation=activation,
+ **tfm_kwargs,
+ )
+ model.fit(series, epochs=1)
+ else:
+ model = base_model(
+ input_chunk_length=1,
+ output_chunk_length=1,
+ activation=activation,
+ **tfm_kwargs,
+ )
+ model.fit(series, epochs=1)
+
+ def test_time_batch_norm_3d(self):
+ torch.manual_seed(0)
+
+ layer = TimeBatchNorm2d()
+ # 4D does not work
+ with pytest.raises(ValueError):
+ layer.forward(torch.randn(3, 3, 3, 3))
+
+ # 2D does not work
+ with pytest.raises(ValueError):
+ layer.forward(torch.randn(3, 3))
+
+ # 3D works
+ norm = layer.forward(torch.randn(3, 3, 3)).detach()
+ assert norm.mean().numpy() == pytest.approx(0.0, abs=0.1)
+ assert norm.std().numpy() == pytest.approx(1.0, abs=0.1)
+
+ @pytest.mark.parametrize("batch_size", [1, 2, 5, 10])
+ def test_time_batch_norm_2d_different_batch_sizes(self, batch_size):
+ layer = TimeBatchNorm2d()
+ input_tensor = torch.randn(batch_size, 3, 3)
+ output = layer.forward(input_tensor)
+ assert output.shape == input_tensor.shape
+
+ def test_time_batch_norm_2d_gradients(self):
+ normalized_shape = (10, 32)
+ layer = TimeBatchNorm2d(normalized_shape)
+ input_tensor = torch.randn(5, 10, 32, requires_grad=True)
+
+ output = layer.forward(input_tensor)
+ output.mean().backward()
+
+ assert input_tensor.grad is not None
diff --git a/docs/source/examples.rst b/docs/source/examples.rst
index 72b2557920..9fd96c177a 100644
--- a/docs/source/examples.rst
+++ b/docs/source/examples.rst
@@ -177,6 +177,16 @@ TiDE model example notebook:
examples/18-TiDE-examples.ipynb
+TimeSeries Mixer (TSMixer) Model
+=======================================
+
+TSMixer model example notebook:
+
+.. toctree::
+ :maxdepth: 1
+
+ 21-TSMixer-examples.ipynb
+
Ensemble Models
=============================
diff --git a/docs/userguide/covariates.md b/docs/userguide/covariates.md
index d9ec6cc72e..cc4c564b87 100644
--- a/docs/userguide/covariates.md
+++ b/docs/userguide/covariates.md
@@ -152,6 +152,7 @@ GFMs are models that can be trained on multiple target (and covariate) time seri
| [DLinearModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.dlinear.html#darts.models.forecasting.dlinear.DLinearModel) | β
| β
| β
|
| [NLinearModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.nlinear.html#darts.models.forecasting.nlinear.NLinearModel) | β
| β
| β
|
| [TiDEModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.tide_model.html#darts.models.forecasting.tide_model.TiDEModel) | β
| β
| β
|
+| [TSMixerModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.tsmixer_model.html#darts.models.forecasting.tsmixer_model.TSMixerModel) | β
| β
| β
|
| Ensemble Models (f) | β
| β
| β
|
**Table 1: Darts' forecasting models and their covariate support**
diff --git a/docs/userguide/torch_forecasting_models.md b/docs/userguide/torch_forecasting_models.md
index 0c2ba84fde..662bc4bc66 100644
--- a/docs/userguide/torch_forecasting_models.md
+++ b/docs/userguide/torch_forecasting_models.md
@@ -116,6 +116,7 @@ Each Torch Forecasting Model inherits from one `{X}CovariatesModel` (covariate c
| `NLinearModel` | | | | | β
|
| `DLinearModel` | | | | | β
|
| `TiDEModel` | | | | | β
|
+| `TSMixerModel` | | | | | β
|
**Table 2: Darts' Torch Forecasting Model covariate support**
diff --git a/examples/21-TSMixer-examples.ipynb b/examples/21-TSMixer-examples.ipynb
new file mode 100644
index 0000000000..1d1735f909
--- /dev/null
+++ b/examples/21-TSMixer-examples.ipynb
@@ -0,0 +1,1025 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Time Series Mixer (TSMixer)\n",
+ "This notebook walks through how to use Darts' `TSMixerModel` and benchmarks it against `TiDEModel`.\n",
+ "\n",
+ "TSMixer (Time-series Mixer) is an all-MLP architecture for time series forecasting. \n",
+ "\n",
+ "It does so by integrating historical time series data, future known inputs, and static contextual information. The architecture uses a combination of conditional feature mixing and mixer layers to process and combine these different types of data for effective forecasting.\n",
+ "\n",
+ "Translated to Darts, this model supports all types of covariates (past, future, and/or static).\n",
+ "\n",
+ "See the original paper and model description [here](https://arxiv.org/abs/2303.06053).\n",
+ "\n",
+ "According to the authors, the model outperforms several state-of-the-art models on multivariate forecasting tasks.\n",
+ "\n",
+ "Let's see how it performs against `TideModel` on the ETTh1 and ETTh2 datasets."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# fix python path if working locally\n",
+ "from utils import fix_pythonpath_if_working_locally\n",
+ "\n",
+ "fix_pythonpath_if_working_locally()\n",
+ "%matplotlib inline"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "%load_ext autoreload\n",
+ "%autoreload 2\n",
+ "%matplotlib inline"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import warnings\n",
+ "\n",
+ "warnings.filterwarnings(\"ignore\")\n",
+ "import logging\n",
+ "\n",
+ "logging.disable(logging.CRITICAL)\n",
+ "\n",
+ "import torch\n",
+ "import matplotlib.pyplot as plt\n",
+ "import numpy as np\n",
+ "import pandas as pd\n",
+ "from pytorch_lightning.callbacks.early_stopping import EarlyStopping\n",
+ "\n",
+ "from darts import concatenate\n",
+ "from darts.dataprocessing.transformers.scaler import Scaler\n",
+ "from darts.datasets import ETTh1Dataset, ETTh2Dataset\n",
+ "from darts.metrics import mae, mse, mql\n",
+ "from darts.models import TiDEModel, TSMixerModel\n",
+ "from darts.utils.likelihood_models import QuantileRegression\n",
+ "from darts.utils.callbacks import TFMProgressBar"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Data Loading and preparation\n",
+ "We consider the ETTh1 and ETTh2 datasets which contain hourly multivariate data of an electricity transformer (load, oil temperature, ...).\n",
+ "You can find more information [here](https://unit8co.github.io/darts/generated_api/darts.datasets.html#darts.datasets.ETTh1Dataset).\n",
+ "\n",
+ "We will add static information to each transformer time series, that identifies whether it is the `ETTh1` or `ETTh2` transformer.\n",
+ "Both TSMixer and TiDE can levarage this information."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " | component | \n",
+ " HUFL | \n",
+ " HULL | \n",
+ " MUFL | \n",
+ " MULL | \n",
+ " LUFL | \n",
+ " LULL | \n",
+ " OT | \n",
+ "
\n",
+ " \n",
+ " | date | \n",
+ " | \n",
+ " | \n",
+ " | \n",
+ " | \n",
+ " | \n",
+ " | \n",
+ " | \n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ " | 2016-07-01 00:00:00 | \n",
+ " 5.827 | \n",
+ " 2.009 | \n",
+ " 1.599 | \n",
+ " 0.462 | \n",
+ " 4.203 | \n",
+ " 1.340 | \n",
+ " 30.531000 | \n",
+ "
\n",
+ " \n",
+ " | 2016-07-01 01:00:00 | \n",
+ " 5.693 | \n",
+ " 2.076 | \n",
+ " 1.492 | \n",
+ " 0.426 | \n",
+ " 4.142 | \n",
+ " 1.371 | \n",
+ " 27.787001 | \n",
+ "
\n",
+ " \n",
+ " | 2016-07-01 02:00:00 | \n",
+ " 5.157 | \n",
+ " 1.741 | \n",
+ " 1.279 | \n",
+ " 0.355 | \n",
+ " 3.777 | \n",
+ " 1.218 | \n",
+ " 27.787001 | \n",
+ "
\n",
+ " \n",
+ " | 2016-07-01 03:00:00 | \n",
+ " 5.090 | \n",
+ " 1.942 | \n",
+ " 1.279 | \n",
+ " 0.391 | \n",
+ " 3.807 | \n",
+ " 1.279 | \n",
+ " 25.044001 | \n",
+ "
\n",
+ " \n",
+ " | 2016-07-01 04:00:00 | \n",
+ " 5.358 | \n",
+ " 1.942 | \n",
+ " 1.492 | \n",
+ " 0.462 | \n",
+ " 3.868 | \n",
+ " 1.279 | \n",
+ " 21.948000 | \n",
+ "
\n",
+ " \n",
+ " | ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ "
\n",
+ " \n",
+ " | 2018-06-26 15:00:00 | \n",
+ " -1.674 | \n",
+ " 3.550 | \n",
+ " -5.615 | \n",
+ " 2.132 | \n",
+ " 3.472 | \n",
+ " 1.523 | \n",
+ " 10.904000 | \n",
+ "
\n",
+ " \n",
+ " | 2018-06-26 16:00:00 | \n",
+ " -5.492 | \n",
+ " 4.287 | \n",
+ " -9.132 | \n",
+ " 2.274 | \n",
+ " 3.533 | \n",
+ " 1.675 | \n",
+ " 11.044000 | \n",
+ "
\n",
+ " \n",
+ " | 2018-06-26 17:00:00 | \n",
+ " 2.813 | \n",
+ " 3.818 | \n",
+ " -0.817 | \n",
+ " 2.097 | \n",
+ " 3.716 | \n",
+ " 1.523 | \n",
+ " 10.271000 | \n",
+ "
\n",
+ " \n",
+ " | 2018-06-26 18:00:00 | \n",
+ " 9.243 | \n",
+ " 3.818 | \n",
+ " 5.472 | \n",
+ " 2.097 | \n",
+ " 3.655 | \n",
+ " 1.432 | \n",
+ " 9.778000 | \n",
+ "
\n",
+ " \n",
+ " | 2018-06-26 19:00:00 | \n",
+ " 10.114 | \n",
+ " 3.550 | \n",
+ " 6.183 | \n",
+ " 1.564 | \n",
+ " 3.716 | \n",
+ " 1.462 | \n",
+ " 9.567000 | \n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
17420 rows Γ 7 columns
\n",
+ "
"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "show_col = \"HUFL\"\n",
+ "for idx, (train_, val_, test_) in enumerate(zip(train, val, test)):\n",
+ " train_[show_col].plot(label=f\"train_trafo_{idx}\")\n",
+ " val_[show_col].plot(label=f\"val_trafo_{idx}\")\n",
+ " test_[show_col].plot(label=f\"test_trafo_{idx}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Now let's scale the data. To avoid leaking information from the validation and test sets, we scale the data based on the properties of the train set."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 7,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "scaler = Scaler() # default uses sklearn's MinMaxScaler\n",
+ "train = scaler.fit_transform(train)\n",
+ "val = scaler.transform(val)\n",
+ "test = scaler.transform(test)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Model Parameter Setup\n",
+ "Boilerplate code is no fun, especially in the context of training multiple models to compare performance. To avoid this, we use a common configuration that can be used with any Darts `TorchForecastingModel`.\n",
+ "\n",
+ "A few interesting things about these parameters:\n",
+ "\n",
+ "- **Gradient clipping:** Mitigates exploding gradients during backpropagation by setting an upper limit on the gradient for a batch.\n",
+ "\n",
+ "- **Learning rate:** The majority of the learning done by a model is in the earlier epochs. As training goes on it is often helpful to reduce the learning rate to fine-tune the model. That being said, it can also lead to significant overfitting.\n",
+ "\n",
+ "- **Early stopping:** To avoid overfitting, we can use early stopping. It monitors a metric on the validation set and stops training once the metric is not improving anymore based on a custom condition.\n",
+ "\n",
+ "- **Likelihood and Loss Functions:** You can either make the model probabilistic with a `likelihood`, or deterministic with a `loss_fn`. In this notebook we train probabilistic models using QuantileRegression.\n",
+ "\n",
+ "- **Reversible Instance Normalization:** Use [Reversible Instance Normalization](https://openreview.net/forum?id=cGDAkQo1C0p) which in most of the cases improves model performance.\n",
+ "\n",
+ "- **Encoders:** We can encode time axis/calendar information and use them as past or future covariates using `add_encoders`. Here, we'll add cyclic encodings of the hour, day of the week, and month as future covariates"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 8,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def create_params(\n",
+ " input_chunk_length: int,\n",
+ " output_chunk_length: int,\n",
+ " full_training=True,\n",
+ "):\n",
+ " # early stopping: this setting stops training once the the validation\n",
+ " # loss has not decreased by more than 1e-5 for 10 epochs\n",
+ " early_stopper = EarlyStopping(\n",
+ " monitor=\"val_loss\",\n",
+ " patience=10,\n",
+ " min_delta=1e-5,\n",
+ " mode=\"min\",\n",
+ " )\n",
+ "\n",
+ " # PyTorch Lightning Trainer arguments (you can add any custom callback)\n",
+ " if full_training:\n",
+ " limit_train_batches = None\n",
+ " limit_val_batches = None\n",
+ " max_epochs = 200\n",
+ " batch_size = 256\n",
+ " else:\n",
+ " limit_train_batches = 20\n",
+ " limit_val_batches = 10\n",
+ " max_epochs = 40\n",
+ " batch_size = 64\n",
+ "\n",
+ " # only show the training and prediction progress bars\n",
+ " progress_bar = TFMProgressBar(\n",
+ " enable_sanity_check_bar=False, enable_validation_bar=False\n",
+ " )\n",
+ " pl_trainer_kwargs = {\n",
+ " \"gradient_clip_val\": 1,\n",
+ " \"max_epochs\": max_epochs,\n",
+ " \"limit_train_batches\": limit_train_batches,\n",
+ " \"limit_val_batches\": limit_val_batches,\n",
+ " \"accelerator\": \"auto\",\n",
+ " \"callbacks\": [early_stopper, progress_bar],\n",
+ " }\n",
+ "\n",
+ " # optimizer setup, uses Adam by default\n",
+ " optimizer_cls = torch.optim.Adam\n",
+ " optimizer_kwargs = {\n",
+ " \"lr\": 1e-4,\n",
+ " }\n",
+ "\n",
+ " # learning rate scheduler\n",
+ " lr_scheduler_cls = torch.optim.lr_scheduler.ExponentialLR\n",
+ " lr_scheduler_kwargs = {\"gamma\": 0.999}\n",
+ "\n",
+ " # for probabilistic models, we use quantile regression, and set `loss_fn` to `None`\n",
+ " likelihood = QuantileRegression()\n",
+ " loss_fn = None\n",
+ "\n",
+ " return {\n",
+ " \"input_chunk_length\": input_chunk_length, # lookback window\n",
+ " \"output_chunk_length\": output_chunk_length, # forecast/lookahead window\n",
+ " \"use_reversible_instance_norm\": True,\n",
+ " \"optimizer_kwargs\": optimizer_kwargs,\n",
+ " \"pl_trainer_kwargs\": pl_trainer_kwargs,\n",
+ " \"lr_scheduler_cls\": lr_scheduler_cls,\n",
+ " \"lr_scheduler_kwargs\": lr_scheduler_kwargs,\n",
+ " \"likelihood\": likelihood, # use a `likelihood` for probabilistic forecasts\n",
+ " \"loss_fn\": loss_fn, # use a `loss_fn` for determinsitic model\n",
+ " \"save_checkpoints\": True, # checkpoint to retrieve the best performing model state,\n",
+ " \"force_reset\": True,\n",
+ " \"batch_size\": batch_size,\n",
+ " \"random_state\": 42,\n",
+ " \"add_encoders\": {\n",
+ " \"cyclic\": {\n",
+ " \"future\": [\"hour\", \"dayofweek\", \"month\"]\n",
+ " } # add cyclic time axis encodings as future covariates\n",
+ " },\n",
+ " }"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Model configuration\n",
+ "Let's use the last week of hourly data as lookback window (`input_chunk_length`) and train a probabilistic model to predict the next 24 hours directly (`output_chunk_length`). Additionally, we tell the model to use the static information. To keep the notebook simple, we'll set `full_training=False`. To get even better performance, set `full_training=True`.\n",
+ "\n",
+ "Apart from that, we use our helper function to set up all the common model arguments."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 9,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "input_chunk_length = 7 * 24\n",
+ "output_chunk_length = 24\n",
+ "use_static_covariates = True\n",
+ "full_training = False"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 10,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# create the models\n",
+ "model_tsm = TSMixerModel(\n",
+ " **create_params(\n",
+ " input_chunk_length,\n",
+ " output_chunk_length,\n",
+ " full_training=full_training,\n",
+ " ),\n",
+ " use_static_covariates=use_static_covariates,\n",
+ " model_name=\"tsm\",\n",
+ ")\n",
+ "model_tide = TiDEModel(\n",
+ " **create_params(\n",
+ " input_chunk_length,\n",
+ " output_chunk_length,\n",
+ " full_training=full_training,\n",
+ " ),\n",
+ " use_static_covariates=use_static_covariates,\n",
+ " model_name=\"tide\",\n",
+ ")\n",
+ "models = {\n",
+ " \"TSM\": model_tsm,\n",
+ " \"TiDE\": model_tide,\n",
+ "}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Model Training"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Now let's train all of the models. When using early stopping it is important to save checkpoints. This allows us to continue past the best model configuration and then restore the optimal weights once training has been completed."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 11,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "application/vnd.jupyter.widget-view+json": {
+ "model_id": "1ab2f4e3c6a14b4687d70b402b9920ac",
+ "version_major": 2,
+ "version_minor": 0
+ },
+ "text/plain": [
+ "Training: 0it [00:00, ?it/s]"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "application/vnd.jupyter.widget-view+json": {
+ "model_id": "c8efee5bcaef467499408860f691509d",
+ "version_major": 2,
+ "version_minor": 0
+ },
+ "text/plain": [
+ "Training: 0it [00:00, ?it/s]"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "# train the models and load the model from its best state/checkpoint\n",
+ "for model_name, model in models.items():\n",
+ " model.fit(\n",
+ " series=train,\n",
+ " val_series=val,\n",
+ " )\n",
+ " # load from checkpoint returns a new model object, we store it in the models dict\n",
+ " models[model_name] = model.load_from_checkpoint(\n",
+ " model_name=model.model_name, best=True\n",
+ " )"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Backtest the probabilistic models\n",
+ "\n",
+ "Let's configure the prediction. For this example, we will:\n",
+ "- generate **historical forecasts** on the test set using the **pre-trained models**. Each forecast covers a 24 hour horizon, and the time between two consecutive forecasts is also 24 hours. This will give us **276 multivariate forecasts per transformer** to evaluate the model!\n",
+ "- generate **500 stochastic samples** for each prediction point (since we have trained probabilistic models)\n",
+ "- evaluate/**backtest** the probabilistic historical forecasts for some quantiles **using the Mean Quantile Loss** (`mql()`).\n",
+ "\n",
+ "And we'll create some helper functions to generating the forecasts, computing the backtest, and to visualize the predictions."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 12,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# configure the probabilistic prediction\n",
+ "num_samples = 500\n",
+ "forecast_horizon = output_chunk_length\n",
+ "\n",
+ "# compute the Mean Quantile Loss over these quantiles\n",
+ "evaluate_quantiles = [0.05, 0.1, 0.2, 0.5, 0.8, 0.9, 0.95]\n",
+ "\n",
+ "\n",
+ "def historical_forecasts(model):\n",
+ " \"\"\"Generates probabilistic historical forecasts for each transformer\n",
+ " and returns the inverse transforms results.\n",
+ "\n",
+ " Each forecast covers 24h (forecast_horizon). The time between two forecasts\n",
+ " (stride) is also 24 hours.\n",
+ " \"\"\"\n",
+ " hfc = model.historical_forecasts(\n",
+ " series=test,\n",
+ " forecast_horizon=forecast_horizon,\n",
+ " stride=forecast_horizon,\n",
+ " last_points_only=False,\n",
+ " retrain=False,\n",
+ " num_samples=num_samples,\n",
+ " verbose=True,\n",
+ " )\n",
+ " return scaler.inverse_transform(hfc)\n",
+ "\n",
+ "\n",
+ "def backtest(model, hfc, name):\n",
+ " \"\"\"Evaluates probabilistic historical forecasts using the Mean Quantile\n",
+ " Loss (MQL) over a set of quantiles.\"\"\"\n",
+ " # add metric specific kwargs\n",
+ " metric_kwargs = [{\"q\": q} for q in evaluate_quantiles]\n",
+ " metrics = [mql for _ in range(len(evaluate_quantiles))]\n",
+ " bt = model.backtest(\n",
+ " series=series,\n",
+ " historical_forecasts=hfc,\n",
+ " last_points_only=False,\n",
+ " metric=metrics,\n",
+ " metric_kwargs=metric_kwargs,\n",
+ " verbose=True,\n",
+ " )\n",
+ " bt = pd.DataFrame(\n",
+ " bt,\n",
+ " columns=[f\"q_{q}\" for q in evaluate_quantiles],\n",
+ " index=[f\"{trafo}_{name}\" for trafo in [\"ETTh1\", \"ETTh2\"]],\n",
+ " )\n",
+ " return bt\n",
+ "\n",
+ "\n",
+ "def generate_plots(n_days, hfcs):\n",
+ " \"\"\"Plot the probabilistic forecasts for each model, transformer and transformer\n",
+ " feature against the ground truth.\"\"\"\n",
+ " # concatenate historical forecasts into contiguous time series\n",
+ " # (works because forecast_horizon=stride)\n",
+ " hfcs_plot = {}\n",
+ " for model_name, hfc_model in hfcs.items():\n",
+ " hfcs_plot[model_name] = [\n",
+ " concatenate(hfc_series[-n_days:], axis=0) for hfc_series in hfc_model\n",
+ " ]\n",
+ "\n",
+ " # remember start and end points for plotting the target series\n",
+ " hfc_ = hfcs_plot[model_name][0]\n",
+ " start, end = hfc_.start_time(), hfc_.end_time()\n",
+ "\n",
+ " # for each target column...\n",
+ " for col in series[0].columns:\n",
+ " fig, axes = plt.subplots(ncols=2, figsize=(12, 6))\n",
+ " # ... and for each transformer...\n",
+ " for trafo_idx, trafo in enumerate(series):\n",
+ " trafo[col][start:end].plot(label=\"ground truth\", ax=axes[trafo_idx])\n",
+ " # ... plot the historical forecasts for each model\n",
+ " for model_name, hfc in hfcs_plot.items():\n",
+ " hfc[trafo_idx][col].plot(\n",
+ " label=model_name + \"_q0.05-q0.95\", ax=axes[trafo_idx]\n",
+ " )\n",
+ " axes[trafo_idx].set_title(f\"ETTh{trafo_idx + 1}: {col}\")\n",
+ " plt.show()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Okay, now we're ready to evaluate the models"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 13,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Model: TSM\n",
+ "Generating historical forecasts..\n"
+ ]
+ },
+ {
+ "data": {
+ "application/vnd.jupyter.widget-view+json": {
+ "model_id": "809ff39dfd7b4192b102d9151b2c1417",
+ "version_major": 2,
+ "version_minor": 0
+ },
+ "text/plain": [
+ "Predicting: 0it [00:00, ?it/s]"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Evaluating historical forecasts..\n",
+ "Model: TiDE\n",
+ "Generating historical forecasts..\n"
+ ]
+ },
+ {
+ "data": {
+ "application/vnd.jupyter.widget-view+json": {
+ "model_id": "ca2e5b2a7d634d7ea619998ce8a11dd7",
+ "version_major": 2,
+ "version_minor": 0
+ },
+ "text/plain": [
+ "Predicting: 0it [00:00, ?it/s]"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Evaluating historical forecasts..\n"
+ ]
+ }
+ ],
+ "source": [
+ "bts = {}\n",
+ "hfcs = {}\n",
+ "for model_name, model in models.items():\n",
+ " print(f\"Model: {model_name}\")\n",
+ " print(\"Generating historical forecasts..\")\n",
+ " hfcs[model_name] = historical_forecasts(models[model_name])\n",
+ "\n",
+ " print(\"Evaluating historical forecasts..\")\n",
+ " bts[model_name] = backtest(models[model_name], hfcs[model_name], model_name)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Let's see how they performed.\n",
+ "\n",
+ "> **Note:** These results are likely to improve/change when setting `full_training=True`"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 14,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " | \n",
+ " q_0.05 | \n",
+ " q_0.1 | \n",
+ " q_0.2 | \n",
+ " q_0.5 | \n",
+ " q_0.8 | \n",
+ " q_0.9 | \n",
+ " q_0.95 | \n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ " | ETTh1_TSM | \n",
+ " 0.501772 | \n",
+ " 0.769545 | \n",
+ " 1.136141 | \n",
+ " 1.568439 | \n",
+ " 1.098847 | \n",
+ " 0.721835 | \n",
+ " 0.442062 | \n",
+ "
\n",
+ " \n",
+ " | ETTh1_TiDE | \n",
+ " 0.573716 | \n",
+ " 0.885452 | \n",
+ " 1.298672 | \n",
+ " 1.671870 | \n",
+ " 1.151501 | \n",
+ " 0.727515 | \n",
+ " 0.446724 | \n",
+ "
\n",
+ " \n",
+ " | ETTh2_TSM | \n",
+ " 0.659187 | \n",
+ " 1.030655 | \n",
+ " 1.508628 | \n",
+ " 1.932923 | \n",
+ " 1.317960 | \n",
+ " 0.857147 | \n",
+ " 0.524620 | \n",
+ "
\n",
+ " \n",
+ " | ETTh2_TiDE | \n",
+ " 0.627251 | \n",
+ " 0.982114 | \n",
+ " 1.450893 | \n",
+ " 1.897117 | \n",
+ " 1.323661 | \n",
+ " 0.862239 | \n",
+ " 0.528638 | \n",
+ "
\n",
+ " \n",
+ "
\n",
+ "