Merge branch 'master' into docs-prototype

CHANGELOG.md

@@ -12,7 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add extension with models from `statsforecast`: `StatsForecastARIMAModel`, `StatsForecastAutoARIMAModel`, `StatsForecastAutoCESModel`, `StatsForecastAutoETSModel`, `StatsForecastAutoThetaModel` ([#1295](https://github.com/tinkoff-ai/etna/pull/1297))
 - Notebook `feature_selection` ([#875](https://github.com/tinkoff-ai/etna/pull/875))
 - 
-- 
+- Implementation of PatchTS model ([#1277](https://github.com/tinkoff-ai/etna/pull/1277))
 
 ### Changed
 - 

etna/models/nn/__init__.py

@@ -6,6 +6,7 @@
  from etna.models.nn.mlp import MLPModel
  from etna.models.nn.nbeats import NBeatsGenericModel
  from etna.models.nn.nbeats import NBeatsInterpretableModel
+ from etna.models.nn.patchts import PatchTSModel
  from etna.models.nn.rnn import RNNModel
  from etna.models.nn.tft import TFTModel
  from etna.models.nn.utils import PytorchForecastingDatasetBuilder

etna/models/nn/patchts.py

@@ -0,0 +1,375 @@
+import math
+from typing import Any
+from typing import Dict
+from typing import Iterator
+from typing import Optional
+
+import numpy as np
+import pandas as pd
+from typing_extensions import TypedDict
+
+from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import FloatDistribution
+from etna.distributions import IntDistribution
+from etna.models.base import DeepBaseModel
+from etna.models.base import DeepBaseNet
+
+if SETTINGS.torch_required:
+ import torch
+ import torch.nn as nn
+
+
+class PatchTSBatch(TypedDict):
+ """Batch specification for PatchTS."""
+
+ encoder_real: "torch.Tensor"
+ decoder_real: "torch.Tensor"
+ encoder_target: "torch.Tensor"
+ decoder_target: "torch.Tensor"
+ segment: "torch.Tensor"
+
+
+class PositionalEncoding(nn.Module):
+
+ """Positional encoding of tokens and reshaping."""
+
+ def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000):
+ super().__init__()
+ self.dropout = nn.Dropout(p=dropout)
+
+ position = torch.arange(max_len).unsqueeze(1)
+ div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
+ pe = torch.zeros(max_len, 1, d_model)
+ pe[:, 0, 0::2] = torch.sin(position * div_term)
+ pe[:, 0, 1::2] = torch.cos(position * div_term)
+ self.register_buffer("pe", pe)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """x: Tensor, shape [batch_size, input_size, patch_num, embedding_dim]."""
+ x = torch.reshape(x, (x.shape[0] * x.shape[1], x.shape[2], x.shape[3]))
+ # x.shape == (batch_size * input_size, patch_num, embedding_dim)
+ x = x.permute(1, 0, 2) # (patch_num, batch_size * input_size, embedding_dim)
+ x = x + self.pe[: x.size(0)]
+ return self.dropout(x)
+
+
+class PatchTSNet(DeepBaseNet):
+ """PatchTS based Lightning module."""
+
+ def __init__(
+ self,
+ encoder_length: int,
+ patch_len: int,
+ stride: int,
+ num_layers: int,
+ hidden_size: int,
+ feedforward_size: int,
+ nhead: int,
+ lr: float,
+ loss: "torch.nn.Module",
+ optimizer_params: Optional[dict],
+ ) -> None:
+ """Init PatchTS.
+
+ Parameters
+ ----------
+ encoder_length:
+ encoder length
+ patch_len:
+ size of patch
+ stride:
+ step of patch
+ num_layers:
+ number of layers
+ hidden_size:
+ size of the hidden state
+ feedforward_size:
+ size of feedforward layers in transformer
+ nhead:
+ number of transformer heads
+ lr:
+ learning rate
+ loss:
+ loss function
+ optimizer_params:
+ parameters for optimizer for Adam optimizer (api reference :py:class:`torch.optim.Adam`)
+ """
+ super().__init__()
+ self.patch_len = patch_len
+ self.num_layers = num_layers
+ self.feedforward_size = feedforward_size
+ self.hidden_size = hidden_size
+ self.nhead = nhead
+ self.stride = stride
+ self.loss = loss
+
+ encoder_layers = nn.TransformerEncoderLayer(
+ d_model=self.hidden_size, nhead=self.nhead, dim_feedforward=self.feedforward_size
+ )
+ self.model = nn.Sequential(
+ nn.Linear(self.patch_len, self.hidden_size),
+ PositionalEncoding(d_model=self.hidden_size),
+ nn.TransformerEncoder(encoder_layers, self.num_layers),
+ )
+ self.max_patch_num = (encoder_length - self.patch_len) // self.stride + 1
+ self.projection = nn.Sequential(
+ nn.Flatten(start_dim=-2), nn.Linear(in_features=self.hidden_size * self.max_patch_num, out_features=1)
+ )
+
+ self.lr = lr
+ self.optimizer_params = {} if optimizer_params is None else optimizer_params
+
+ def forward(self, x: PatchTSBatch, *args, **kwargs): # type: ignore
+ """Forward pass.
+
+ Parameters
+ ----------
+ x:
+ batch of data
+
+ Returns
+ -------
+ :
+ forecast with shape (batch_size, decoder_length, 1)
+ """
+ encoder_real = x["encoder_real"].float() # (batch_size, encoder_length, input_size)
+ decoder_real = x["decoder_real"].float() # (batch_size, decoder_length, input_size)
+ decoder_length = decoder_real.shape[1]
+ outputs = []
+ current_input = encoder_real
+ for _ in range(decoder_length):
+ pred = self._get_prediction(current_input)
+ outputs.append(pred)
+ current_input = torch.cat((current_input[:, 1:, :], torch.unsqueeze(pred, dim=1)), dim=1)
+
+ forecast = torch.cat(outputs, dim=1)
+ forecast = torch.unsqueeze(forecast, dim=2)
+
+ return forecast
+
+ def _get_prediction(self, x: torch.Tensor) -> torch.Tensor:
+ x = x.permute(0, 2, 1) # (batch_size, input_size, encoder_length)
+ # do patching
+ x = x.unfold(
+ dimension=-1, size=self.patch_len, step=self.stride
+ ) # (batch_size, input_size, patch_num, patch_len)
+
+ y = self.model(x)
+ y = y.permute(1, 0, 2) # (batch_size, hidden_size, patch_num)
+
+ return self.projection(y) # (batch_size, 1)
+
+ def step(self, batch: PatchTSBatch, *args, **kwargs): # type: ignore
+ """Step for loss computation for training or validation.
+
+ Parameters
+ ----------
+ batch:
+ batch of data
+
+ Returns
+ -------
+ :
+ loss, true_target, prediction_target
+ """
+ encoder_real = batch["encoder_real"].float() # (batch_size, encoder_length, input_size)
+ decoder_real = batch["decoder_real"].float() # (batch_size, decoder_length, input_size)
+
+ decoder_target = batch["decoder_target"].float() # (batch_size, decoder_length, 1)
+
+ decoder_length = decoder_real.shape[1]
+
+ outputs = []
+ x = encoder_real
+ for i in range(decoder_length):
+ pred = self._get_prediction(x)
+ outputs.append(pred)
+ x = torch.cat((x[:, 1:, :], torch.unsqueeze(decoder_real[:, i, :], dim=1)), dim=1)
+
+ target_prediction = torch.cat(outputs, dim=1)
+ target_prediction = torch.unsqueeze(target_prediction, dim=2)
+
+ loss = self.loss(target_prediction, decoder_target)
+ return loss, decoder_target, target_prediction
+
+ def make_samples(self, df: pd.DataFrame, encoder_length: int, decoder_length: int) -> Iterator[dict]:
+ """Make samples from segment DataFrame."""
+ values_real = df.select_dtypes(include=[np.number]).values
+ values_target = df["target"].values
+ segment = df["segment"].values[0]
+
+ def _make(
+ values_real: np.ndarray,
+ values_target: np.ndarray,
+ segment: str,
+ start_idx: int,
+ encoder_length: int,
+ decoder_length: int,
+ ) -> Optional[dict]:
+
+ sample: Dict[str, Any] = {
+ "encoder_real": list(),
+ "decoder_real": list(),
+ "encoder_target": list(),
+ "decoder_target": list(),
+ "segment": None,
+ }
+ total_length = len(values_target)
+ total_sample_length = encoder_length + decoder_length
+
+ if total_sample_length + start_idx > total_length:
+ return None
+
+ sample["decoder_real"] = values_real[start_idx + encoder_length : start_idx + total_sample_length]
+ sample["encoder_real"] = values_real[start_idx : start_idx + encoder_length]
+
+ target = values_target[start_idx : start_idx + encoder_length + decoder_length].reshape(-1, 1)
+ sample["encoder_target"] = target[:encoder_length]
+ sample["decoder_target"] = target[encoder_length:]
+
+ sample["segment"] = segment
+
+ return sample
+
+ start_idx = 0
+ while True:
+ batch = _make(
+ values_target=values_target,
+ values_real=values_real,
+ segment=segment,
+ start_idx=start_idx,
+ encoder_length=encoder_length,
+ decoder_length=decoder_length,
+ )
+ if batch is None:
+ break
+ yield batch
+ start_idx += 1
+
+ def configure_optimizers(self) -> "torch.optim.Optimizer":
+ """Optimizer configuration."""
+ optimizer = torch.optim.Adam(self.parameters(), lr=self.lr, **self.optimizer_params)
+ return optimizer
+
+
+class PatchTSModel(DeepBaseModel):
+ """PatchTS model using PyTorch layers."""
+
+ def __init__(
+ self,
+ decoder_length: int,
+ encoder_length: int,
+ patch_len: int = 4,
+ stride: int = 1,
+ num_layers: int = 3,
+ hidden_size: int = 128,
+ feedforward_size: int = 256,
+ nhead: int = 16,
+ lr: float = 1e-3,
+ loss: Optional["torch.nn.Module"] = None,
+ train_batch_size: int = 128,
+ test_batch_size: int = 128,
+ optimizer_params: Optional[dict] = None,
+ trainer_params: Optional[dict] = None,
+ train_dataloader_params: Optional[dict] = None,
+ test_dataloader_params: Optional[dict] = None,
+ val_dataloader_params: Optional[dict] = None,
+ split_params: Optional[dict] = None,
+ ):
+ """Init PatchTS model.
+
+ Parameters
+ ----------
+ encoder_length:
+ encoder length
+ decoder_length:
+ decoder length
+ patch_len:
+ size of patch
+ stride:
+ step of patch
+ num_layers:
+ number of layers
+ hidden_size:
+ size of the hidden state
+ feedforward_size:
+ size of feedforward layers in transformer
+ nhead:
+ number of transformer heads
+ lr:
+ learning rate
+ loss:
+ loss function, MSELoss by default
+ train_batch_size:
+ batch size for training
+ test_batch_size:
+ batch size for testing
+ optimizer_params:
+ parameters for optimizer for Adam optimizer (api reference :py:class:`torch.optim.Adam`)
+ trainer_params:
+ Pytorch ligthning trainer parameters (api reference :py:class:`pytorch_lightning.trainer.trainer.Trainer`)
+ train_dataloader_params:
+ parameters for train dataloader like sampler for example (api reference :py:class:`torch.utils.data.DataLoader`)
+ test_dataloader_params:
+ parameters for test dataloader
+ val_dataloader_params:
+ parameters for validation dataloader
+ split_params:
+ dictionary with parameters for :py:func:`torch.utils.data.random_split` for train-test splitting
+ * **train_size**: (*float*) value from 0 to 1 - fraction of samples to use for training
+
+ * **generator**: (*Optional[torch.Generator]*) - generator for reproducibile train-test splitting
+
+ * **torch_dataset_size**: (*Optional[int]*) - number of samples in dataset, in case of dataset not implementing ``__len__``
+ """
+ self.num_layers = num_layers
+ self.hidden_size = hidden_size
+ self.lr = lr
+ self.patch_len = patch_len
+ self.stride = stride
+ self.nhead = nhead
+ self.feedforward_size = feedforward_size
+ self.loss = loss if loss is not None else nn.MSELoss()
+ self.optimizer_params = optimizer_params
+ super().__init__(
+ net=PatchTSNet(
+ encoder_length,
+ patch_len=self.patch_len,
+ stride=self.stride,
+ num_layers=self.num_layers,
+ hidden_size=self.hidden_size,
+ feedforward_size=self.feedforward_size,
+ nhead=self.nhead,
+ lr=self.lr,
+ loss=self.loss,
+ optimizer_params=self.optimizer_params,
+ ),
+ decoder_length=decoder_length,
+ encoder_length=encoder_length,
+ train_batch_size=train_batch_size,
+ test_batch_size=test_batch_size,
+ train_dataloader_params=train_dataloader_params,
+ test_dataloader_params=test_dataloader_params,
+ val_dataloader_params=val_dataloader_params,
+ trainer_params=trainer_params,
+ split_params=split_params,
+ )
+
+ def params_to_tune(self) -> Dict[str, BaseDistribution]:
+ """Get default grid for tuning hyperparameters.
+
+ This grid tunes parameters: ``num_layers``, ``hidden_size``, ``lr``, ``encoder_length``.
+ Other parameters are expected to be set by the user.
+
+ Returns
+ -------
+ :
+ Grid to tune.
+ """
+ return {
+ "num_layers": IntDistribution(low=1, high=3),
+ "hidden_size": IntDistribution(low=16, high=256, step=self.nhead),
+ "lr": FloatDistribution(low=1e-5, high=1e-2, log=True),
+ "encoder_length": IntDistribution(low=self.patch_len, high=24),
+ }