Copy_of_Transformers_for_timeseries.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1tQfKhP0-OOp"
      },
      "source": [
        "## Transformers for Timeseries\n",
        "\n",
        "Click to run on colab (if you're not already there): [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/charlesollion/dlexperiments/blob/master/7-Transformers-Timeseries/Transformers_for_timeseries.ipynb)\n",
        "\n",
        "The goal of this notebook is to illustrate the use of a transformer for timeseries prediction.\n",
        "This notebook was built by Alice Martin and adapted to pytorch by Charles Ollion"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "qrOtpRDjrSY2"
      },
      "outputs": [
        {
          "ename": "ValueError",
          "evalue": "libcublas.so.*[0-9] not found in the system path ['/home/wcyat/dev/weather-predict', '/home/wcyat/.pyenv/versions/3.10.12/lib/python310.zip', '/home/wcyat/.pyenv/versions/3.10.12/lib/python3.10', '/home/wcyat/.pyenv/versions/3.10.12/lib/python3.10/lib-dynload', '', '/home/wcyat/dev/weather-predict/.venv/lib/python3.10/site-packages']",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mOSError\u001b[0m                                   Traceback (most recent call last)",
            "File \u001b[0;32m~/dev/weather-predict/.venv/lib/python3.10/site-packages/torch/__init__.py:168\u001b[0m, in \u001b[0;36m_load_global_deps\u001b[0;34m()\u001b[0m\n\u001b[1;32m    167\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> 168\u001b[0m     ctypes\u001b[39m.\u001b[39;49mCDLL(lib_path, mode\u001b[39m=\u001b[39;49mctypes\u001b[39m.\u001b[39;49mRTLD_GLOBAL)\n\u001b[1;32m    169\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mOSError\u001b[39;00m \u001b[39mas\u001b[39;00m err:\n\u001b[1;32m    170\u001b[0m     \u001b[39m# Can only happen for wheel with cuda libs as PYPI deps\u001b[39;00m\n\u001b[1;32m    171\u001b[0m     \u001b[39m# As PyTorch is not purelib, but nvidia-*-cu11 is\u001b[39;00m\n",
            "File \u001b[0;32m~/.pyenv/versions/3.10.12/lib/python3.10/ctypes/__init__.py:374\u001b[0m, in \u001b[0;36mCDLL.__init__\u001b[0;34m(self, name, mode, handle, use_errno, use_last_error, winmode)\u001b[0m\n\u001b[1;32m    373\u001b[0m \u001b[39mif\u001b[39;00m handle \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[0;32m--> 374\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_handle \u001b[39m=\u001b[39m _dlopen(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_name, mode)\n\u001b[1;32m    375\u001b[0m \u001b[39melse\u001b[39;00m:\n",
            "\u001b[0;31mOSError\u001b[0m: libcurand.so.10: cannot open shared object file: No such file or directory",
            "\nDuring handling of the above exception, another exception occurred:\n",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[3], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mtorch\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mnumpy\u001b[39;00m \u001b[39mas\u001b[39;00m \u001b[39mnp\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mpandas\u001b[39;00m \u001b[39mas\u001b[39;00m \u001b[39mpd\u001b[39;00m\n",
            "File \u001b[0;32m~/dev/weather-predict/.venv/lib/python3.10/site-packages/torch/__init__.py:228\u001b[0m\n\u001b[1;32m    217\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m    218\u001b[0m     \u001b[39m# Easy way.  You want this most of the time, because it will prevent\u001b[39;00m\n\u001b[1;32m    219\u001b[0m     \u001b[39m# C++ symbols from libtorch clobbering C++ symbols from other\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    225\u001b[0m     \u001b[39m#\u001b[39;00m\n\u001b[1;32m    226\u001b[0m     \u001b[39m# See Note [Global dependencies]\u001b[39;00m\n\u001b[1;32m    227\u001b[0m     \u001b[39mif\u001b[39;00m USE_GLOBAL_DEPS:\n\u001b[0;32m--> 228\u001b[0m         _load_global_deps()\n\u001b[1;32m    229\u001b[0m     \u001b[39mfrom\u001b[39;00m \u001b[39mtorch\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39m_C\u001b[39;00m \u001b[39mimport\u001b[39;00m \u001b[39m*\u001b[39m  \u001b[39m# noqa: F403\u001b[39;00m\n\u001b[1;32m    231\u001b[0m \u001b[39m# Appease the type checker; ordinarily this binding is inserted by the\u001b[39;00m\n\u001b[1;32m    232\u001b[0m \u001b[39m# torch._C module initialization code in C\u001b[39;00m\n",
            "File \u001b[0;32m~/dev/weather-predict/.venv/lib/python3.10/site-packages/torch/__init__.py:189\u001b[0m, in \u001b[0;36m_load_global_deps\u001b[0;34m()\u001b[0m\n\u001b[1;32m    187\u001b[0m     \u001b[39mraise\u001b[39;00m err\n\u001b[1;32m    188\u001b[0m \u001b[39mfor\u001b[39;00m lib_folder, lib_name \u001b[39min\u001b[39;00m cuda_libs\u001b[39m.\u001b[39mitems():\n\u001b[0;32m--> 189\u001b[0m     _preload_cuda_deps(lib_folder, lib_name)\n\u001b[1;32m    190\u001b[0m ctypes\u001b[39m.\u001b[39mCDLL(lib_path, mode\u001b[39m=\u001b[39mctypes\u001b[39m.\u001b[39mRTLD_GLOBAL)\n",
            "File \u001b[0;32m~/dev/weather-predict/.venv/lib/python3.10/site-packages/torch/__init__.py:154\u001b[0m, in \u001b[0;36m_preload_cuda_deps\u001b[0;34m(lib_folder, lib_name)\u001b[0m\n\u001b[1;32m    152\u001b[0m         \u001b[39mbreak\u001b[39;00m\n\u001b[1;32m    153\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m lib_path:\n\u001b[0;32m--> 154\u001b[0m     \u001b[39mraise\u001b[39;00m \u001b[39mValueError\u001b[39;00m(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m{\u001b[39;00mlib_name\u001b[39m}\u001b[39;00m\u001b[39m not found in the system path \u001b[39m\u001b[39m{\u001b[39;00msys\u001b[39m.\u001b[39mpath\u001b[39m}\u001b[39;00m\u001b[39m\"\u001b[39m)\n\u001b[1;32m    155\u001b[0m ctypes\u001b[39m.\u001b[39mCDLL(lib_path)\n",
            "\u001b[0;31mValueError\u001b[0m: libcublas.so.*[0-9] not found in the system path ['/home/wcyat/dev/weather-predict', '/home/wcyat/.pyenv/versions/3.10.12/lib/python310.zip', '/home/wcyat/.pyenv/versions/3.10.12/lib/python3.10', '/home/wcyat/.pyenv/versions/3.10.12/lib/python3.10/lib-dynload', '', '/home/wcyat/dev/weather-predict/.venv/lib/python3.10/site-packages']"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from sklearn.model_selection import train_test_split\n",
        "import time\n",
        "import matplotlib.pyplot as plt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hMUQEguW7dj4"
      },
      "source": [
        "### Preparing the Dataset\n",
        "Energy consumption dataset from https://archive.ics.uci.edu/ml/datasets/Appliances+energy+prediction\n",
        "* gathers 10-min measurements of household appliances energy consumption (20 first features), coupled with local meteorological data (8 last features).\n",
        "* The time-series forecasting task is to predict the first 20 features, given as input data the 28 features. A window of observations of 12 time steps is considered to predict the next series of observations (this corresponds to a 2-hours window of observations.\n",
        "\n",
        "you may get the dataset (a single csv file) by running the following cell:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z4X_aQcF7A8m"
      },
      "outputs": [],
      "source": [
        "!wget https://raw.githubusercontent.com/LuisM78/Appliances-energy-prediction-data/master/energydata_complete.csv"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C7-ULhqewgMi"
      },
      "outputs": [],
      "source": [
        "# load and preprocess the energy dataset:\n",
        "def convert_col_into_float(df, list_cols):\n",
        "    for col in list_cols:\n",
        "        df[col] = df[col].astype(str)\n",
        "        df[col] = df[col].str.replace(',', '.')\n",
        "        df[col] = df[col].astype(np.float32)\n",
        "    return df\n",
        "df = pd.read_csv(\"energydata_complete.csv\", index_col='date', parse_dates=['date'])\n",
        "print(df.head())\n",
        "list_cols = list(df.columns)\n",
        "# gathers 10-min measurements of household appliances energy consumption (20 first features), coupled with local meteorological data. (8 last features)\n",
        "print(\"dataset variables\", list_cols)\n",
        "df = convert_col_into_float(df, list_cols)\n",
        "data = df.values"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p0bQB8uOv7xT"
      },
      "outputs": [],
      "source": [
        "def split_dataset_into_seq(dataset, start_index=0, end_index=None, history_size=13, step=1):\n",
        "    '''split the dataset to have sequence of observations of length history size'''\n",
        "    data = []\n",
        "    start_index = start_index + history_size\n",
        "    if end_index is None:\n",
        "        end_index = len(dataset)\n",
        "    for i in range(start_index, end_index):\n",
        "        indices = range(i - history_size, i, step)\n",
        "        data.append(dataset[indices])\n",
        "    return np.array(data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ajb5bl5bvtkh"
      },
      "outputs": [],
      "source": [
        "def split_dataset(data, TRAIN_SPLIT=0.7, VAL_SPLIT=0.5, save_path=None):\n",
        "    '''split the dataset into train, val and test splits'''\n",
        "    # normalization\n",
        "    data_mean = data.mean(axis=0)\n",
        "    data_std = data.std(axis=0)\n",
        "    data = (data - data_mean) / data_std\n",
        "    stats = (data_mean, data_std)\n",
        "\n",
        "    data_in_seq = split_dataset_into_seq(data, start_index=0, end_index=None, history_size=13, step=1)\n",
        "\n",
        "    # split between validation dataset and test set:\n",
        "    train_data, val_data = train_test_split(data_in_seq, train_size=TRAIN_SPLIT, shuffle=True, random_state=123)\n",
        "    val_data, test_data = train_test_split(val_data, train_size=VAL_SPLIT, shuffle=True, random_state=123)\n",
        "\n",
        "    return train_data, val_data, test_data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "34lFfEqjzdlB"
      },
      "outputs": [],
      "source": [
        "def split_fn(chunk):\n",
        "    \"\"\"to split the dataset sequences into input and targets sequences\"\"\"\n",
        "    inputs = torch.tensor(chunk[:, :-1, :], device=device)\n",
        "    targets = torch.tensor(chunk[:, 1:, :], device=device)\n",
        "    return inputs, targets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n1ClBYcty4jL"
      },
      "outputs": [],
      "source": [
        "def data_to_dataset(train_data, val_data, test_data, batch_size=32, target_features=list(range(20))):\n",
        "    '''\n",
        "    split each train split into inputs and targets\n",
        "    convert each train split into a tf.dataset\n",
        "    '''\n",
        "    x_train, y_train = split_fn(train_data)\n",
        "    x_val, y_val = split_fn(val_data)\n",
        "    x_test, y_test = split_fn(test_data)\n",
        "    # selecting only the first 20 features for prediction:\n",
        "    y_train = y_train[:, :, target_features]\n",
        "    y_val = y_val[:, :, target_features]\n",
        "    y_test = y_test[:, :, target_features]\n",
        "    train_dataset = torch.utils.data.TensorDataset(x_train, y_train)\n",
        "    val_dataset = torch.utils.data.TensorDataset(x_val, y_val)\n",
        "    test_dataset = torch.utils.data.TensorDataset(x_test, y_test)\n",
        "\n",
        "    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size)\n",
        "    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size)\n",
        "    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size)\n",
        "    return train_loader, val_loader, test_loader"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pUsnkjlK-OOp"
      },
      "outputs": [],
      "source": [
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VDzKvost8VUu"
      },
      "outputs": [],
      "source": [
        "train_data, val_data, test_data = split_dataset(data)\n",
        "train_dataset, val_dataset, test_dataset = data_to_dataset(train_data, val_data, test_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tGacA_fS85sz"
      },
      "source": [
        "### Implementation of the Transformer model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8KH-wFTB4e7I"
      },
      "outputs": [],
      "source": [
        "import torch.nn as nn\n",
        "import math\n",
        "\n",
        "class MultiHeadAttention(nn.Module):\n",
        "    '''Multi-head self-attention module'''\n",
        "    def __init__(self, D, H):\n",
        "        super(MultiHeadAttention, self).__init__()\n",
        "        self.H = H # number of heads\n",
        "        self.D = D # dimension\n",
        "\n",
        "        self.wq = nn.Linear(D, D*H)\n",
        "        self.wk = nn.Linear(D, D*H)\n",
        "        self.wv = nn.Linear(D, D*H)\n",
        "\n",
        "        self.dense = nn.Linear(D*H, D)\n",
        "\n",
        "    def concat_heads(self, x):\n",
        "        '''(B, H, S, D) => (B, S, D*H)'''\n",
        "        B, H, S, D = x.shape\n",
        "        x = x.permute((0, 2, 1, 3)).contiguous()  # (B, S, H, D)\n",
        "        x = x.reshape((B, S, H*D))   # (B, S, D*H)\n",
        "        return x\n",
        "\n",
        "    def split_heads(self, x):\n",
        "        '''(B, S, D*H) => (B, H, S, D)'''\n",
        "        B, S, D_H = x.shape\n",
        "        x = x.reshape(B, S, self.H, self.D)    # (B, S, H, D)\n",
        "        x = x.permute((0, 2, 1, 3))  # (B, H, S, D)\n",
        "        return x\n",
        "\n",
        "    def forward(self, x, mask):\n",
        "\n",
        "        q = self.wq(x)  # (B, S, D*H)\n",
        "        k = self.wk(x)  # (B, S, D*H)\n",
        "        v = self.wv(x)  # (B, S, D*H)\n",
        "\n",
        "        q = self.split_heads(q)  # (B, H, S, D)\n",
        "        k = self.split_heads(k)  # (B, H, S, D)\n",
        "        v = self.split_heads(v)  # (B, H, S, D)\n",
        "\n",
        "        attention_scores = torch.matmul(q, k.transpose(-1, -2)) #(B,H,S,S)\n",
        "        attention_scores = attention_scores / math.sqrt(self.D)\n",
        "\n",
        "        # add the mask to the scaled tensor.\n",
        "        if mask is not None:\n",
        "            attention_scores += (mask * -1e9)\n",
        "\n",
        "        attention_weights = nn.Softmax(dim=-1)(attention_scores)\n",
        "        scaled_attention = torch.matmul(attention_weights, v)  # (B, H, S, D)\n",
        "        concat_attention = self.concat_heads(scaled_attention) # (B, S, D*H)\n",
        "        output = self.dense(concat_attention)  # (B, S, D)\n",
        "\n",
        "        return output, attention_weights"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ADmmyqoU-OOp"
      },
      "outputs": [],
      "source": [
        "B, S, H, D = 9, 11, 5, 8\n",
        "mha = MultiHeadAttention(D, H)\n",
        "out, att = mha.forward(torch.zeros(B, S, D), mask=None)\n",
        "out.shape, att.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bsiqoBF15qxb"
      },
      "outputs": [],
      "source": [
        "# Positional encodings\n",
        "def get_angles(pos, i, D):\n",
        "    angle_rates = 1 / np.power(10000, (2 * (i // 2)) / np.float32(D))\n",
        "    return pos * angle_rates\n",
        "\n",
        "\n",
        "def positional_encoding(D, position=20, dim=3, device=device):\n",
        "    angle_rads = get_angles(np.arange(position)[:, np.newaxis],\n",
        "                            np.arange(D)[np.newaxis, :],\n",
        "                            D)\n",
        "    # apply sin to even indices in the array; 2i\n",
        "    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])\n",
        "    # apply cos to odd indices in the array; 2i+1\n",
        "    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])\n",
        "    if dim == 3:\n",
        "        pos_encoding = angle_rads[np.newaxis, ...]\n",
        "    elif dim == 4:\n",
        "        pos_encoding = angle_rads[np.newaxis,np.newaxis,  ...]\n",
        "    return torch.tensor(pos_encoding, device=device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZW26NwrB6LtR"
      },
      "outputs": [],
      "source": [
        "# function that implement the look_ahead mask for masking future time steps.\n",
        "def create_look_ahead_mask(size, device=device):\n",
        "    mask = torch.ones((size, size), device=device)\n",
        "    mask = torch.triu(mask, diagonal=1)\n",
        "    return mask  # (size, size)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_MGvkIo6-OOp"
      },
      "outputs": [],
      "source": [
        "create_look_ahead_mask(6)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KUalYf-dtYwb"
      },
      "outputs": [],
      "source": [
        "class TransformerLayer(nn.Module):\n",
        "    def __init__(self, D, H, hidden_mlp_dim, dropout_rate):\n",
        "        super(TransformerLayer, self).__init__()\n",
        "        self.dropout_rate = dropout_rate\n",
        "        self.mlp_hidden = nn.Linear(D, hidden_mlp_dim)\n",
        "        self.mlp_out = nn.Linear(hidden_mlp_dim, D)\n",
        "        self.layernorm1 = nn.LayerNorm(D, eps=1e-9)\n",
        "        self.layernorm2 = nn.LayerNorm(D, eps=1e-9)\n",
        "        self.dropout1 = nn.Dropout(dropout_rate)\n",
        "        self.dropout2 = nn.Dropout(dropout_rate)\n",
        "\n",
        "        self.mha = MultiHeadAttention(D, H)\n",
        "\n",
        "\n",
        "    def forward(self, x, look_ahead_mask):\n",
        "\n",
        "        attn, attn_weights = self.mha(x, look_ahead_mask)  # (B, S, D)\n",
        "        attn = self.dropout1(attn) # (B,S,D)\n",
        "        attn = self.layernorm1(attn + x) # (B,S,D)\n",
        "\n",
        "        mlp_act = torch.relu(self.mlp_hidden(attn))\n",
        "        mlp_act = self.mlp_out(mlp_act)\n",
        "        mlp_act = self.dropout2(mlp_act)\n",
        "\n",
        "        output = self.layernorm2(mlp_act + attn)  # (B, S, D)\n",
        "\n",
        "        return output, attn_weights"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P1UBGsft-OOp"
      },
      "outputs": [],
      "source": [
        "dl = TransformerLayer(16, 3, 32, 0.1)\n",
        "out, attn = dl(x=torch.zeros(5, 7, 16), look_ahead_mask=None)\n",
        "out.shape, attn.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "04VLiWfcuD4d"
      },
      "outputs": [],
      "source": [
        "class Transformer(nn.Module):\n",
        "    '''Transformer Decoder Implementating several Decoder Layers.\n",
        "    '''\n",
        "    def __init__(self, num_layers, D, H, hidden_mlp_dim, inp_features, out_features, dropout_rate):\n",
        "        super(Transformer, self).__init__()\n",
        "        self.sqrt_D = torch.tensor(math.sqrt(D))\n",
        "        self.num_layers = num_layers\n",
        "        self.input_projection = nn.Linear(inp_features, D) # multivariate input\n",
        "        self.output_projection = nn.Linear(D, out_features) # multivariate output\n",
        "        self.pos_encoding = positional_encoding(D)\n",
        "        self.dec_layers = nn.ModuleList([TransformerLayer(D, H, hidden_mlp_dim,\n",
        "                                        dropout_rate=dropout_rate\n",
        "                                       ) for _ in range(num_layers)])\n",
        "        self.dropout = nn.Dropout(dropout_rate)\n",
        "\n",
        "    def forward(self, x, mask):\n",
        "        B, S, D = x.shape\n",
        "        attention_weights = {}\n",
        "        x = self.input_projection(x)\n",
        "        x *= self.sqrt_D\n",
        "\n",
        "        x += self.pos_encoding[:, :S, :]\n",
        "\n",
        "        x = self.dropout(x)\n",
        "\n",
        "        for i in range(self.num_layers):\n",
        "            x, block = self.dec_layers[i](x=x,\n",
        "                                          look_ahead_mask=mask)\n",
        "            attention_weights['decoder_layer{}'.format(i + 1)] = block\n",
        "\n",
        "        x = self.output_projection(x)\n",
        "\n",
        "        return x, attention_weights # (B,S,S)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qD8WG0B--V1x"
      },
      "outputs": [],
      "source": [
        "# Test Forward pass on the Transformer:\n",
        "transformer = Transformer(num_layers=1, D=32, H=1, hidden_mlp_dim=32,\n",
        "                                       inp_features=28, out_features=20, dropout_rate=0.1)\n",
        "transformer.to(device)\n",
        "(inputs, targets) = next(iter(train_dataset))\n",
        "\n",
        "S = inputs.shape[1]\n",
        "mask = create_look_ahead_mask(S)\n",
        "out, attn = transformer (x=inputs, mask=mask)\n",
        "out.shape, attn[\"decoder_layer1\"].shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WzrZUFZz6ia_"
      },
      "source": [
        "## Training the Transformer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IEfyfebM-OOp"
      },
      "outputs": [],
      "source": [
        "param_sizes = [p.numel() for p in transformer.parameters()]\n",
        "print(f\"number of weight/biases matrices: {len(param_sizes)} \"\n",
        "      f\"for a total of {np.sum(param_sizes)} parameters \")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bMGE5tSk-OOp"
      },
      "outputs": [],
      "source": [
        "transformer = Transformer(num_layers=1, D=32, H=4, hidden_mlp_dim=32,\n",
        "                          inp_features=28, out_features=20, dropout_rate=0.1).to(device)\n",
        "optimizer = torch.optim.RMSprop(transformer.parameters(),\n",
        "                                lr=0.00005)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nkwUseV6-OOp"
      },
      "outputs": [],
      "source": [
        "from tqdm import tqdm\n",
        "\n",
        "n_epochs = 20\n",
        "niter = len(train_dataset)\n",
        "losses, val_losses = [], []\n",
        "\n",
        "for e in tqdm(range(n_epochs)):\n",
        "\n",
        "    # one epoch on train set\n",
        "    transformer.train()\n",
        "    sum_train_loss = 0.0\n",
        "    for x,y in train_dataset:\n",
        "        S = x.shape[1]\n",
        "        mask = create_look_ahead_mask(S)\n",
        "        out, _ = transformer(x, mask)\n",
        "        loss = torch.nn.MSELoss()(out, y)\n",
        "        sum_train_loss += loss.item()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "    losses.append(sum_train_loss / niter)\n",
        "\n",
        "    # Evaluate on val set\n",
        "    transformer.eval()\n",
        "    sum_val_loss = 0.0\n",
        "    for i, (x, y) in enumerate(val_dataset):\n",
        "        S = x.shape[1]\n",
        "        mask = create_look_ahead_mask(S)\n",
        "        out, _ = transformer(x, mask)\n",
        "        loss = torch.nn.MSELoss()(out, y)\n",
        "        sum_val_loss += loss.item()\n",
        "    val_losses.append(sum_val_loss / (i + 1))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iw71pT_t-OOp"
      },
      "outputs": [],
      "source": [
        "plt.plot(losses)\n",
        "plt.plot(val_losses);"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ALeWAQDFnAGx"
      },
      "source": [
        "### Evaluation on Test set"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zRJJ-39R257V"
      },
      "outputs": [],
      "source": [
        "test_losses, test_preds  = [], []\n",
        "transformer.eval()\n",
        "for (x, y) in test_dataset:\n",
        "    S = x.shape[-2]\n",
        "    y_pred, _ = transformer(x,\n",
        "                            mask=create_look_ahead_mask(S))\n",
        "    loss_test = torch.nn.MSELoss()(y_pred, y)  # (B,S)\n",
        "    test_losses.append(loss_test.item())\n",
        "    test_preds.append(y_pred.detach().cpu().numpy())\n",
        "test_preds = np.vstack(test_preds)\n",
        "np.mean(test_losses)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k1VCvxaklm0k"
      },
      "outputs": [],
      "source": [
        "# Display predictions vs ground truth:\n",
        "# we'll take one random element of the first batch\n",
        "# and display the first feature\n",
        "seq_len = 12\n",
        "index = np.random.randint(32)\n",
        "feature_num = 0\n",
        "\n",
        "x_test, _ = test_dataset.dataset.tensors\n",
        "x_test = x_test[index, :, feature_num].cpu().numpy()\n",
        "pred = test_preds[index, :, feature_num]\n",
        "x = np.linspace(1, seq_len, seq_len)\n",
        "plt.plot(x, pred, 'red', lw=2, label='predictions for sample: {}'.format(index))\n",
        "plt.plot(x, x_test, 'cyan', lw=2, label='ground-truth for sample: {}'.format(index))\n",
        "plt.legend(fontsize=10)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pirSKNR9-OOp"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Copy of Transformers-for-timeseries.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.12"
    },
    "toc": {
      "base_numbering": 1,
      "nav_menu": {},
      "number_sections": true,
      "sideBar": true,
      "skip_h1_title": false,
      "title_cell": "Table of Contents",
      "title_sidebar": "Contents",
      "toc_cell": false,
      "toc_position": {},
      "toc_section_display": true,
      "toc_window_display": false
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}