models.itransformer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp models.itransformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "from fastcore.test import test_eq\n",
    "from nbdev.showdoc import show_doc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# iTransformer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The iTransformer model simply takes the Transformer architecture but it applies the attention and feed-forward network on the inverted dimensions. This means that time points of each individual series are embedded into tokens. That way, the attention mechanisms learn multivariate correlation and the feed-forward network learns non-linear relationships.\n",
    "\n",
    "**References**\n",
    "- [Yong Liu, Tengge Hu, Haoran Zhang, Haixu Wu, Shiyu Wang, Lintao Ma, Mingsheng Long. \"iTransformer: Inverted Transformers Are Effective for Time Series Forecasting\"](https://arxiv.org/abs/2310.06625)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Figure 1. Architecture of iTransformer.](imgs_models/itransformer.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from typing import Optional\n",
    "from math import sqrt\n",
    "\n",
    "from neuralforecast.losses.pytorch import MAE\n",
    "from neuralforecast.common._base_multivariate import BaseMultivariate\n",
    "\n",
    "from neuralforecast.common._modules import TransEncoder, TransEncoderLayer, AttentionLayer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Auxiliary functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.1 Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "\n",
    "class TriangularCausalMask():\n",
    "    def __init__(self, B, L, device=\"cpu\"):\n",
    "        mask_shape = [B, 1, L, L]\n",
    "        with torch.no_grad():\n",
    "            self._mask = torch.triu(torch.ones(mask_shape, dtype=torch.bool), diagonal=1).to(device)\n",
    "\n",
    "    @property\n",
    "    def mask(self):\n",
    "        return self._mask\n",
    "\n",
    "class FullAttention(nn.Module):\n",
    "    def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):\n",
    "        super(FullAttention, self).__init__()\n",
    "        self.scale = scale\n",
    "        self.mask_flag = mask_flag\n",
    "        self.output_attention = output_attention\n",
    "        self.dropout = nn.Dropout(attention_dropout)\n",
    "\n",
    "    def forward(self, queries, keys, values, attn_mask, tau=None, delta=None):\n",
    "        B, L, H, E = queries.shape\n",
    "        _, S, _, D = values.shape\n",
    "        scale = self.scale or 1. / sqrt(E)\n",
    "\n",
    "        scores = torch.einsum(\"blhe,bshe->bhls\", queries, keys)\n",
    "\n",
    "        if self.mask_flag:\n",
    "            if attn_mask is None:\n",
    "                attn_mask = TriangularCausalMask(B, L, device=queries.device)\n",
    "\n",
    "            scores.masked_fill_(attn_mask.mask, -np.inf)\n",
    "\n",
    "        A = self.dropout(torch.softmax(scale * scores, dim=-1))\n",
    "        V = torch.einsum(\"bhls,bshd->blhd\", A, values)\n",
    "\n",
    "        if self.output_attention:\n",
    "            return (V.contiguous(), A)\n",
    "        else:\n",
    "            return (V.contiguous(), None)      "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.2 Inverted embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "\n",
    "class DataEmbedding_inverted(nn.Module):\n",
    "    def __init__(self, c_in, hidden_size, dropout=0.1):\n",
    "        super(DataEmbedding_inverted, self).__init__()\n",
    "        self.value_embedding = nn.Linear(c_in, hidden_size)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "\n",
    "    def forward(self, x, x_mark):\n",
    "        x = x.permute(0, 2, 1)\n",
    "        # x: [Batch Variate Time]\n",
    "        if x_mark is None:\n",
    "            x = self.value_embedding(x)\n",
    "        else:\n",
    "            # the potential to take covariates (e.g. timestamps) as tokens\n",
    "            x = self.value_embedding(torch.cat([x, x_mark.permute(0, 2, 1)], 1)) \n",
    "        # x: [Batch Variate hidden_size]\n",
    "        return self.dropout(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "\n",
    "class iTransformer(BaseMultivariate):\n",
    "\n",
    "    \"\"\" iTransformer\n",
    "\n",
    "    **Parameters:**<br>\n",
    "    `h`: int, Forecast horizon. <br>\n",
    "    `input_size`: int, autorregresive inputs size, y=[1,2,3,4] input_size=2 -> y_[t-2:t]=[1,2].<br>\n",
    "    `n_series`: int, number of time-series.<br>\n",
    "    `futr_exog_list`: str list, future exogenous columns.<br>\n",
    "    `hist_exog_list`: str list, historic exogenous columns.<br>\n",
    "    `stat_exog_list`: str list, static exogenous columns.<br>\n",
    "    `hidden_size`: int, dimension of the model.<br>\n",
    "    `n_heads`: int, number of heads.<br>\n",
    "    `e_layers`: int, number of encoder layers.<br>\n",
    "    `d_layers`: int, number of decoder layers.<br>\n",
    "    `d_ff`: int, dimension of fully-connected layer.<br>\n",
    "    `factor`: int, attention factor.<br>\n",
    "    `dropout`: float, dropout rate.<br>\n",
    "    `use_norm`: bool, whether to normalize or not.<br>\n",
    "    `loss`: PyTorch module, instantiated train loss class from [losses collection](https://nixtla.github.io/neuralforecast/losses.pytorch.html).<br>\n",
    "    `valid_loss`: PyTorch module=`loss`, instantiated valid loss class from [losses collection](https://nixtla.github.io/neuralforecast/losses.pytorch.html).<br>\n",
    "    `max_steps`: int=1000, maximum number of training steps.<br>\n",
    "    `learning_rate`: float=1e-3, Learning rate between (0, 1).<br>\n",
    "    `num_lr_decays`: int=-1, Number of learning rate decays, evenly distributed across max_steps.<br>\n",
    "    `early_stop_patience_steps`: int=-1, Number of validation iterations before early stopping.<br>\n",
    "    `val_check_steps`: int=100, Number of training steps between every validation loss check.<br>\n",
    "    `batch_size`: int=32, number of different series in each batch.<br>\n",
    "    `step_size`: int=1, step size between each window of temporal data.<br>\n",
    "    `scaler_type`: str='identity', type of scaler for temporal inputs normalization see [temporal scalers](https://nixtla.github.io/neuralforecast/common.scalers.html).<br>\n",
    "    `random_seed`: int=1, random_seed for pytorch initializer and numpy generators.<br>\n",
    "    `num_workers_loader`: int=os.cpu_count(), workers to be used by `TimeSeriesDataLoader`.<br>\n",
    "    `drop_last_loader`: bool=False, if True `TimeSeriesDataLoader` drops last non-full batch.<br>\n",
    "    `alias`: str, optional,  Custom name of the model.<br>\n",
    "    `optimizer`: Subclass of 'torch.optim.Optimizer', optional, user specified optimizer instead of the default choice (Adam).<br>\n",
    "    `optimizer_kwargs`: dict, optional, list of parameters used by the user specified `optimizer`.<br>\n",
    "    `**trainer_kwargs`: int,  keyword trainer arguments inherited from [PyTorch Lighning's trainer](https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.trainer.trainer.Trainer.html?highlight=trainer).<br>\n",
    "    \n",
    "    **References**<br>\n",
    "    - [Yong Liu, Tengge Hu, Haoran Zhang, Haixu Wu, Shiyu Wang, Lintao Ma, Mingsheng Long. \"iTransformer: Inverted Transformers Are Effective for Time Series Forecasting\"](https://arxiv.org/abs/2310.06625)\n",
    "    \"\"\"\n",
    "\n",
    "    # Class attributes\n",
    "    SAMPLING_TYPE = 'multivariate'\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 input_size,\n",
    "                 n_series,\n",
    "                 futr_exog_list = None,\n",
    "                 hist_exog_list = None,\n",
    "                 stat_exog_list = None,\n",
    "                 hidden_size: int = 512,\n",
    "                 n_heads: int = 8,\n",
    "                 e_layers: int = 2,\n",
    "                 d_layers: int = 1,\n",
    "                 d_ff: int = 2048,\n",
    "                 factor: int = 1,\n",
    "                 dropout: float = 0.1,\n",
    "                 use_norm: bool = True,\n",
    "                 loss = MAE(),\n",
    "                 valid_loss = None,\n",
    "                 max_steps: int = 1000,\n",
    "                 learning_rate: float = 1e-3,\n",
    "                 num_lr_decays: int = -1,\n",
    "                 early_stop_patience_steps: int =-1,\n",
    "                 val_check_steps: int = 100,\n",
    "                 batch_size: int = 32,\n",
    "                 step_size: int = 1,\n",
    "                 scaler_type: str = 'identity',\n",
    "                 random_seed: int = 1,\n",
    "                 num_workers_loader: int = 0,\n",
    "                 drop_last_loader: bool = False,\n",
    "                 optimizer = None,\n",
    "                 optimizer_kwargs = None,\n",
    "                 **trainer_kwargs):\n",
    "        \n",
    "        super(iTransformer, self).__init__(h=h,\n",
    "                                           input_size=input_size,\n",
    "                                           n_series=n_series,\n",
    "                                           stat_exog_list = None,\n",
    "                                           futr_exog_list = None,\n",
    "                                           hist_exog_list = None,\n",
    "                                           loss=loss,\n",
    "                                           valid_loss=valid_loss,\n",
    "                                           max_steps=max_steps,\n",
    "                                           learning_rate=learning_rate,\n",
    "                                           num_lr_decays=num_lr_decays,\n",
    "                                           early_stop_patience_steps=early_stop_patience_steps,\n",
    "                                           val_check_steps=val_check_steps,\n",
    "                                           batch_size=batch_size,\n",
    "                                           step_size=step_size,\n",
    "                                           scaler_type=scaler_type,\n",
    "                                           random_seed=random_seed,\n",
    "                                           num_workers_loader=num_workers_loader,\n",
    "                                           drop_last_loader=drop_last_loader,\n",
    "                                           optimizer=optimizer,\n",
    "                                           optimizer_kwargs=optimizer_kwargs,\n",
    "                                           **trainer_kwargs)\n",
    "        \n",
    "        # Asserts\n",
    "        if stat_exog_list is not None:\n",
    "            raise Exception(\"iTransformer does not support static exogenous variables\")\n",
    "        if futr_exog_list is not None:\n",
    "            raise Exception(\"iTransformer does not support future exogenous variables\")\n",
    "        if hist_exog_list is not None:\n",
    "            raise Exception(\"iTransformer does not support historical exogenous variables\")\n",
    "        \n",
    "        self.enc_in = n_series\n",
    "        self.dec_in = n_series\n",
    "        self.c_out = n_series\n",
    "        self.hidden_size = hidden_size\n",
    "        self.n_heads = n_heads\n",
    "        self.e_layers = e_layers\n",
    "        self.d_layers = d_layers\n",
    "        self.d_ff = d_ff\n",
    "        self.factor = factor\n",
    "        self.dropout = dropout\n",
    "        self.use_norm = use_norm\n",
    "\n",
    "        # Architecture\n",
    "        self.enc_embedding = DataEmbedding_inverted(input_size, self.hidden_size, self.dropout)\n",
    "\n",
    "        self.encoder = TransEncoder(\n",
    "            [\n",
    "                TransEncoderLayer(\n",
    "                    AttentionLayer(\n",
    "                        FullAttention(False, self.factor, attention_dropout=self.dropout), self.hidden_size, self.n_heads),\n",
    "                    self.hidden_size,\n",
    "                    self.d_ff,\n",
    "                    dropout=self.dropout,\n",
    "                    activation=F.gelu\n",
    "                ) for l in range(self.e_layers)\n",
    "            ],\n",
    "            norm_layer=torch.nn.LayerNorm(self.hidden_size)\n",
    "        )\n",
    "\n",
    "        self.projector = nn.Linear(self.hidden_size, h, bias=True)\n",
    "    \n",
    "    def forecast(self, x_enc):\n",
    "        if self.use_norm:\n",
    "            # Normalization from Non-stationary Transformer\n",
    "            means = x_enc.mean(1, keepdim=True).detach()\n",
    "            x_enc = x_enc - means\n",
    "            stdev = torch.sqrt(torch.var(x_enc, dim=1, keepdim=True, unbiased=False) + 1e-5)\n",
    "            x_enc /= stdev\n",
    "\n",
    "        _, _, N = x_enc.shape # B L N\n",
    "        # B: batch_size;       E: hidden_size; \n",
    "        # L: input_size;       S: horizon(h);\n",
    "        # N: number of variate (tokens), can also includes covariates\n",
    "\n",
    "        # Embedding\n",
    "        # B L N -> B N E                (B L N -> B L E in the vanilla Transformer)\n",
    "        enc_out = self.enc_embedding(x_enc, None) # covariates (e.g timestamp) can be also embedded as tokens\n",
    "        \n",
    "        # B N E -> B N E                (B L E -> B L E in the vanilla Transformer)\n",
    "        # the dimensions of embedded time series has been inverted, and then processed by native attn, layernorm and ffn modules\n",
    "        enc_out, attns = self.encoder(enc_out, attn_mask=None)\n",
    "\n",
    "        # B N E -> B N S -> B S N \n",
    "        dec_out = self.projector(enc_out).permute(0, 2, 1)[:, :, :N] # filter the covariates\n",
    "\n",
    "        if self.use_norm:\n",
    "            # De-Normalization from Non-stationary Transformer\n",
    "            dec_out = dec_out * (stdev[:, 0, :].unsqueeze(1).repeat(1, self.h, 1))\n",
    "            dec_out = dec_out + (means[:, 0, :].unsqueeze(1).repeat(1, self.h, 1))\n",
    "\n",
    "        return dec_out\n",
    "    \n",
    "    def forward(self, windows_batch):\n",
    "        insample_y = windows_batch['insample_y']\n",
    "\n",
    "        y_pred = self.forecast(insample_y)\n",
    "        y_pred = y_pred[:, -self.h:, :]\n",
    "        y_pred = self.loss.domain_map(y_pred)\n",
    "\n",
    "        # domain_map might have squeezed the last dimension in case n_series == 1\n",
    "        if y_pred.ndim == 2:\n",
    "            return y_pred.unsqueeze(-1)\n",
    "        else:\n",
    "            return y_pred\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(iTransformer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(iTransformer.fit, name='iTransformer.fit')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(iTransformer.predict, name='iTransformer.predict')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Usage example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pytorch_lightning as pl\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from neuralforecast import NeuralForecast\n",
    "from neuralforecast.utils import AirPassengersPanel, AirPassengersStatic\n",
    "from neuralforecast.losses.pytorch import MSE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y_train_df = AirPassengersPanel[AirPassengersPanel.ds<AirPassengersPanel['ds'].values[-12]].reset_index(drop=True) # 132 train\n",
    "Y_test_df = AirPassengersPanel[AirPassengersPanel.ds>=AirPassengersPanel['ds'].values[-12]].reset_index(drop=True) # 12 test\n",
    "\n",
    "model = iTransformer(h=12,\n",
    "                     input_size=24,\n",
    "                     n_series=2,\n",
    "                     hidden_size=128,\n",
    "                     n_heads=2,\n",
    "                     e_layers=2,\n",
    "                     d_layers=1,\n",
    "                     d_ff=4,\n",
    "                     factor=1,\n",
    "                     dropout=0.1,\n",
    "                     use_norm=True,\n",
    "                     loss=MSE(),\n",
    "                     valid_loss=MAE(),\n",
    "                     early_stop_patience_steps=3,\n",
    "                     batch_size=32)\n",
    "\n",
    "fcst = NeuralForecast(models=[model], freq='M')\n",
    "fcst.fit(df=Y_train_df, static_df=AirPassengersStatic, val_size=12)\n",
    "forecasts = fcst.predict(futr_df=Y_test_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| eval: false\n",
    "# Plot predictions\n",
    "fig, ax = plt.subplots(1, 1, figsize = (20, 7))\n",
    "Y_hat_df = forecasts.reset_index(drop=False).drop(columns=['unique_id','ds'])\n",
    "plot_df = pd.concat([Y_test_df, Y_hat_df], axis=1)\n",
    "plot_df = pd.concat([Y_train_df, plot_df])\n",
    "\n",
    "plot_df = plot_df[plot_df.unique_id=='Airline1'].drop('unique_id', axis=1)\n",
    "plt.plot(plot_df['ds'], plot_df['y'], c='black', label='True')\n",
    "plt.plot(plot_df['ds'], plot_df['iTransformer'], c='blue', label='Forecast')\n",
    "ax.set_title('AirPassengers Forecast', fontsize=22)\n",
    "ax.set_ylabel('Monthly Passengers', fontsize=20)\n",
    "ax.set_xlabel('Year', fontsize=20)\n",
    "ax.legend(prop={'size': 15})\n",
    "ax.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pytorch_lightning as pl\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from neuralforecast import NeuralForecast\n",
    "from neuralforecast.utils import AirPassengersPanel, AirPassengersStatic\n",
    "from neuralforecast.losses.pytorch import MSE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y_train_df = AirPassengersPanel[AirPassengersPanel.ds<AirPassengersPanel['ds'].values[-12]].reset_index(drop=True) # 132 train\n",
    "Y_test_df = AirPassengersPanel[AirPassengersPanel.ds>=AirPassengersPanel['ds'].values[-12]].reset_index(drop=True) # 12 test\n",
    "\n",
    "model = iTransformer(h=12,\n",
    "                     input_size=24,\n",
    "                     n_series=1,\n",
    "                     hidden_size=128,\n",
    "                     n_heads=2,\n",
    "                     e_layers=2,\n",
    "                     d_layers=1,\n",
    "                     d_ff=4,\n",
    "                     factor=1,\n",
    "                     dropout=0.1,\n",
    "                     use_norm=True,\n",
    "                     loss=MSE(),\n",
    "                     valid_loss=MAE(),\n",
    "                     early_stop_patience_steps=3,\n",
    "                     batch_size=32)\n",
    "\n",
    "fcst = NeuralForecast(models=[model], freq='M')\n",
    "fcst.fit(df=Y_train_df, val_size=12)\n",
    "forecasts = fcst.predict(futr_df=Y_test_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}