experiments_simmtm-BeijingPM25Quality.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center><img src=\"../img/img_0.PNG\"  width=\"1000\" height=\"240\"/></center>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "This notebook demonstrate an apllication of SimMTM, a simple self-supervised learning framework for time series modeling. Self-supervised learning is a learning paradigm that allows model to learn a good representation from the input data itself. The learned representation will be beneficial to some downstream tasks such as forecasting, classification and outlier detection. \n",
    "\n",
    "Self-supervised learning has a lof of success and achieves state-of-the-art performance in some domains, especially in the image domain. In this demo, we will show a self-supervised learning method, SimMTM, in the time-series domain. SimMTM adopts both masked modeling and contrastive modeling to learn a good representation of the input data. By using the learned representation and finetuning it, we achieve a significant improvement compared to the model without self-supervised learning. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/shamvinc/ssl_time_series/mvts_transformer/src\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "try:\n",
    "    os.chdir('src')\n",
    "except:\n",
    "    pass\n",
    "print(os.getcwd())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%reload_ext autoreload\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import copy\n",
    "\n",
    "import math\n",
    "\n",
    "from datasets.datasplit import split_dataset\n",
    "from datasets.data import data_factory, Normalizer, TSRegressionArchive, CSVRegressionArchive\n",
    "from datasets.datasplit import split_dataset\n",
    "from datasets.dataset import collate_superv\n",
    "from models.ts_transformer import model_factory\n",
    "from models.loss import get_loss_module, contrastive_loss\n",
    "from optimizers import get_optimizer\n",
    "\n",
    "from options import Options\n",
    "from running import setup\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Masked Modeling\n",
    "\n",
    "Self-supervision via a ‘pretext task’ on input data combined with finetuning on labeled data is widely used for improving model performance in language and computer\n",
    "vision. One of the popular self-supervision tasks on language data is masked modeling. Masking modeling is to mask some of the input entries randomly and predict those masked entries by using unmasked entries. By masked modeling, the model can learn the relationship through different features and different timesteps. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"../img/img_1.PNG\"  width=\"900\" height=\"240\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"../img/img_2.PNG\"  width=\"900\" height=\"240\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Masking Choice\n",
    "### Random Masking\n",
    "\n",
    "Random Masking is not a good choice to learn a good representation because the model can simply learn to take the average from the neighbour values. \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "<img src=\"../img/img_3.PNG\" width=\"600\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Geometric Masking\n",
    "\n",
    "Instead, we choose to use the geometric masking method, which is to mask a sequence of the input data randomly. The length of the sequence is followed by a geometric distribution. In this case, the model requires to recover a masked sequence from other unmasked input data. We suggest the expected length of a masked sequence is a half of the whole time series sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def geom_noise_mask_single(L, lm, masking_ratio):\n",
    "    \"\"\"\n",
    "    Randomly create a boolean mask of length `L`, consisting of subsequences of average length lm, masking with 0s a `masking_ratio`\n",
    "    proportion of the sequence L. The length of masking subsequences and intervals follow a geometric distribution.\n",
    "    Args:\n",
    "        L: length of mask and sequence to be masked\n",
    "        lm: average length of masking subsequences (streaks of 0s)\n",
    "        masking_ratio: proportion of L to be masked\n",
    "\n",
    "    Returns:\n",
    "        (L,) boolean numpy array intended to mask ('drop') with 0s a sequence of length L\n",
    "    \"\"\"\n",
    "    keep_mask = np.ones(L, dtype=bool)\n",
    "    p_m = 1 / lm  # probability of each masking sequence stopping. parameter of geometric distribution.\n",
    "    p_u = p_m * masking_ratio / (1 - masking_ratio)  # probability of each unmasked sequence stopping. parameter of geometric distribution.\n",
    "    p = [p_m, p_u]\n",
    "\n",
    "    # Start in state 0 with masking_ratio probability\n",
    "    state = int(np.random.rand() > masking_ratio)  # state 0 means masking, 1 means not masking\n",
    "    for i in range(L):\n",
    "        keep_mask[i] = state  # here it happens that state and masking value corresponding to state are identical\n",
    "        if np.random.rand() < p[state]:\n",
    "            state = 1 - state\n",
    "\n",
    "    return keep_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# SimMTM ultilizes both contrastive learning and mask modeling to learn the data representation.\n",
    "## 1 - Contrastive Learning\n",
    "\n",
    "when we mask the input time series data, we create many masked views of the input data. We expect that the distance between two views of the same time series sequence is minimized while maximizing the distance between two different sequences."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "<img src=\"../img/img_5.png\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The contrastive loss is the following: (Eq. 8 in the paper)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "<center><img src=\"../img/img_6.PNG\"/><center/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def demo_contrastive_loss(s, batch_size, tau=0.05):\n",
    "    s = s.squeeze(-1) \n",
    "\n",
    "    B = s.shape[0]\n",
    "    v = s.reshape(B, -1)\n",
    "\n",
    "    norm_v = torch.norm(v, p=2, dim=-1).unsqueeze(-1)\n",
    "    v = v/norm_v\n",
    "    u = torch.transpose(v, 0, 1)\n",
    "\n",
    "    R = torch.matmul(v,u)\n",
    "\n",
    " \n",
    "    R = torch.exp(R/tau) # (batch + mask size) x (batch + mask size)\n",
    "    \n",
    "    # number of masks\n",
    "    M = B//batch_size\n",
    "    mask = torch.eye(batch_size, device=R.device).repeat_interleave(M,dim=0).repeat_interleave(M,dim=1)\n",
    "\n",
    "    denom = R * (torch.ones_like(R) - torch.eye(R.shape[0], device=R.device))\n",
    "\n",
    "    denom = R.sum(-1).unsqueeze(-1)\n",
    "\n",
    "    loss = torch.log(R/denom)\n",
    "    \n",
    "\n",
    "    loss = (loss * (mask - torch.eye(R.shape[0], device=R.device))).sum(1)/(M-1) # except no masked unit\n",
    "    loss = loss.mean(0)\n",
    "    \n",
    "    return -loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2 - Masked Modeling\n",
    "\n",
    "SimMTM proposes to recover a time serie by the weighted sum of multiple masked points, which eases the reconstruction task by assembling ruined but complementary temporal variations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "<img src=\"../img/img_4.png\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from models.ts_transformer import LearnablePositionalEncoding, TransformerBatchNormEncoderLayer\n",
    "\n",
    "class DemoSimMTMTransformerEncoder(nn.Module):\n",
    "    \n",
    "    def __init__(self, max_len, feat_dim, out_len, out_dim, d_model=16, n_heads=4, num_layers=2, dim_feedforward=32, dropout=0.2, temporal_unit=3):\n",
    "        super(DemoSimMTMTransformerEncoder, self).__init__()\n",
    "\n",
    "        self.max_len = max_len\n",
    "        self.d_model = d_model\n",
    "        self.n_heads = n_heads\n",
    "        \n",
    "        self.tau = 0.05\n",
    "        self.mask_length = max_len//2\n",
    "        self.mask_rate = 0.5\n",
    "\n",
    "        self.project_inp = nn.Linear(feat_dim, d_model)\n",
    "        self.projector_layer = nn.Linear(max_len, 1)\n",
    "        self.pos_enc1 = LearnablePositionalEncoding(d_model, dropout=dropout, max_len=max_len)\n",
    "        self.pos_enc2 = LearnablePositionalEncoding(d_model, dropout=dropout, max_len=out_len)\n",
    "        \n",
    "        self.act = F.gelu \n",
    "\n",
    "        # encoder_layer = nn.TransformerEncoderLayer(d_model, self.n_heads, dim_feedforward, dropout, activation='gelu')\n",
    "        encoder_layer = TransformerBatchNormEncoderLayer(d_model, self.n_heads, dim_feedforward, dropout, activation='gelu')\n",
    "\n",
    "        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers)\n",
    "\n",
    "        self.output_layer = nn.Linear(d_model, feat_dim)\n",
    "\n",
    "        self.dropout1 = nn.Dropout(dropout)\n",
    "        self.dropout2 = nn.Dropout(dropout)\n",
    "        self.dropout3 = nn.Dropout2d(dropout)\n",
    "\n",
    "        # self.predict_layer1 = nn.Conv1d(d_model, 512, 5, stride=1)\n",
    "        self.predict_layer1 = nn.Linear(max_len, out_len)\n",
    "        self.predict_layer2 = nn.Linear(d_model, out_dim)\n",
    "        # self.bn = nn.BatchNorm1d(d_model)\n",
    "\n",
    "        self.feat_dim = feat_dim\n",
    "\n",
    "        self.temporal_unit = temporal_unit\n",
    "\n",
    "        self.w1 = torch.nn.parameter.Parameter(data=torch.ones(1), requires_grad=True)\n",
    "        self.w2 = torch.nn.parameter.Parameter(data=torch.ones(1), requires_grad=True)\n",
    "        \n",
    "        \n",
    "    def forward(self, X):\n",
    "        \"\"\"\n",
    "        Reconstruct the input and create the projected output of X\n",
    "        \n",
    "        Args:\n",
    "            X: (batch_size, seq_length, feat_dim) torch tensor of original input\n",
    "\n",
    "        Returns:\n",
    "            output: (batch_size, seq_length, feat_dim)\n",
    "            s: (batch_size, d_model, 1)\n",
    "        \"\"\"\n",
    "\n",
    "        \n",
    "        _x = X\n",
    "         \n",
    "        # Create masked views of the input X\n",
    "        for i in range(self.temporal_unit):\n",
    "            mask = geom_noise_mask_single(X.shape[0] * X.shape[1] * X.shape[2], self.mask_length, self.mask_rate)\n",
    "            mask = mask.reshape(X.shape[0], X.shape[1], X.shape[2])\n",
    "            mask = torch.from_numpy(mask).to(X.device)\n",
    "            x_masked = mask * X\n",
    "            _x = torch.cat([_x, x_masked], axis=-1) # [batch_size, seq_length, feat_dim * temporal_unit]\n",
    "    \n",
    "        \n",
    "        _x = _x.reshape(X.shape[0] * (self.temporal_unit + 1), X.shape[1], X.shape[2])\n",
    "  \n",
    "\n",
    "        inp = _x.permute(1, 0, 2)\n",
    "        inp = self.project_inp(inp) * np.sqrt(self.d_model)  # [seq_length, batch_size, d_model] project input vectors to d_model dimensional space\n",
    "        inp = self.pos_enc1(inp)  # add positional encoding\n",
    "\n",
    "        \n",
    "        output = self.transformer_encoder(inp)  # (seq_length, batch_size, d_model)\n",
    "        output = self.act(output)  # the output transformer encoder/decoder embeddings don't include non-linearity\n",
    "        output = output.permute(1, 0, 2)  # (batch_size, seq_length, d_model)\n",
    "        output = self.dropout1(output)\n",
    "\n",
    "        z_hat, _s = self.project(output, self.tau)\n",
    "        # Most probably defining a Linear(d_model,feat_dim) vectorizes the operation over (seq_length, batch_size).\n",
    "        output = self.output_layer(z_hat)  # (batch_size, seq_length, feat_dim)\n",
    "\n",
    "        return output, _s\n",
    "\n",
    "    \n",
    "    def project(self, z, tau):\n",
    "        \"\"\"\n",
    "        Output a weighted average of z\n",
    "        \n",
    "        Args:\n",
    "            X: (batch_size, seq_length, feat_dim) torch tensor of original input\n",
    "\n",
    "        Returns:\n",
    "            z_hat: (batch_size, seq_length, d_model)\n",
    "            s: (batch_size, d_model, 1)\n",
    "        \"\"\"\n",
    "        _z = z.transpose(1, 2) # [batch_size, d_model, seq_length]\n",
    "        _s = s = self.projector_layer(_z) # [batch_size, d_model, 1]\n",
    "        \n",
    "        if self.training:\n",
    "            mask = torch.ones(1, self.d_model, 1).to(z.device)\n",
    "            mask = self.dropout3(mask)\n",
    "            s = s * mask \n",
    "            s = s + torch.randn(s.shape).to(z.device) * 1e-2\n",
    "        \n",
    "        \n",
    "        s = s.squeeze(-1) \n",
    "        B = s.shape[0]\n",
    "        v = s.reshape(B, -1)\n",
    "\n",
    "        norm_v = torch.norm(v, p=2, dim=-1).unsqueeze(-1)\n",
    "        v = v/norm_v\n",
    "        u = torch.transpose(v, 0, 1)\n",
    "        \n",
    "        R = torch.matmul(v,u)\n",
    "     \n",
    "  \n",
    "        R = torch.exp(R/tau) # (batch + mask size) x (batch + mask size)\n",
    "        R = R * (torch.ones_like(R) - torch.eye(R.shape[0], device=R.device)) # zero out the weight of no masked component\n",
    "        R = R/R.sum(-1).unsqueeze(-1)\n",
    "        M = self.temporal_unit + 1\n",
    "        R = R[::M] # extract every no mask unit # (batch size) x (batch + mask size)\n",
    "\n",
    "        z_hat = (R.unsqueeze(-1).unsqueeze(-1) * z.unsqueeze(0)).sum(1) \n",
    "        return z_hat, _s\n",
    "\n",
    "\n",
    "    def predict(self, X):\n",
    "        \"\"\"\n",
    "        Predict an output given X\n",
    "        \n",
    "        Args:\n",
    "            z: (batch_size, seq_length, d_model) torch tensor of representations of input\n",
    "            tau: temperture of similarity matrix\n",
    "\n",
    "        Returns:\n",
    "            output: (batch_size, out_seq_len, out_dim)\n",
    "        \"\"\"\n",
    "        \n",
    "        # permute because pytorch convention for transformers is [seq_length, batch_size, feat_dim]. padding_masks [batch_size, feat_dim]\n",
    "        inp = X.permute(1, 0, 2)\n",
    "        inp = self.project_inp(inp) * np.sqrt(self.d_model)  # [seq_length, batch_size, d_model] project input vectors to d_model dimensional space\n",
    "        inp = self.pos_enc1(inp)  # add positional encoding\n",
    "        # NOTE: logic for padding masks is reversed to comply with definition in MultiHeadAttention, TransformerEncoderLayer\n",
    "\n",
    "        output = self.transformer_encoder(inp)\n",
    "        output = output.permute(1, 0, 2)  # (batch_size, seq_length, d_model)\n",
    "        # output = self.dropout1(output)\n",
    "       \n",
    "        output = output.transpose(1, 2) # (batch_size, d_model, seq_length)\n",
    "        output = self.predict_layer1(output)\n",
    "        # output = self.act(output)\n",
    "        \n",
    "        output = output.transpose(1, 2) # (batch_size, seq_length, d_model)\n",
    "        output = output.permute(1, 0, 2)\n",
    "        output = self.pos_enc2(output)\n",
    "        output = output.permute(1, 0, 2)\n",
    "        output = self.dropout2(output)\n",
    "        output = self.predict_layer2(output) \n",
    "        return output\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Loading and Preparation\n",
    "\n",
    "In this demo, we use a benchmask time series dataset called BeijingPM25Quality.\n",
    "This dataset is part of the Monash, UEA & UCR time series regression repository. http://tseregression.org/\n",
    "\n",
    "The goal of this dataset is to predict PM2.5 air quality in the city of Beijing. This dataset contains 17532 time series with 9 dimensions.  This includes hourly air pollutants measurments (SO2, NO2, CO and O3), temperature, pressure, dew point, rainfall and windspeed measurments from 12 nationally controlled air quality monitoring sites. The air-quality data are from the Beijing Municipal Environmental Monitoring Center. The meteorological data in each air-quality site are matched with the nearest weather station from the China Meteorological Administration. The time period is from March 1st, 2013 to February 28th, 2017. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-24 07:55:47,195 | INFO : Stored configuration file in '../experiments/_2023-08-24_07-55-46_ETs'\n",
      "11942it [00:47, 251.82it/s]\n",
      "5072it [00:19, 257.36it/s]\n"
     ]
    }
   ],
   "source": [
    "args = Options().parse()  \n",
    "args.data_dir = '../datasets/BeijingPM25Quality'\n",
    "args.task = 'regression'\n",
    "args.output_dir = '../experiments'\n",
    "config = setup(args)\n",
    "from datasets.data import CSVRegressionArchive\n",
    "data = TSRegressionArchive(config['data_dir'], pattern='TRAIN', config=config)\n",
    "test_data = TSRegressionArchive(config['data_dir'], pattern='TEST', config=config)\n",
    "_data = data\n",
    "\n",
    "# Standard Normalization\n",
    "normalizer = Normalizer(config['normalization'])\n",
    "data.feature_df = normalizer.normalize(data.feature_df)\n",
    "test_data.feature_df = normalizer.normalize(test_data.feature_df)\n",
    "\n",
    "\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>dim_0</th>\n",
       "      <th>dim_1</th>\n",
       "      <th>dim_2</th>\n",
       "      <th>dim_3</th>\n",
       "      <th>dim_4</th>\n",
       "      <th>dim_5</th>\n",
       "      <th>dim_6</th>\n",
       "      <th>dim_7</th>\n",
       "      <th>dim_8</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>-0.578722</td>\n",
       "      <td>-1.262337</td>\n",
       "      <td>-0.827452</td>\n",
       "      <td>0.331285</td>\n",
       "      <td>-1.354744</td>\n",
       "      <td>1.282595</td>\n",
       "      <td>-1.671894</td>\n",
       "      <td>-0.077886</td>\n",
       "      <td>2.197098</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>-0.578722</td>\n",
       "      <td>-1.262337</td>\n",
       "      <td>-0.827452</td>\n",
       "      <td>0.331285</td>\n",
       "      <td>-1.390910</td>\n",
       "      <td>1.302248</td>\n",
       "      <td>-1.626941</td>\n",
       "      <td>-0.077886</td>\n",
       "      <td>2.439030</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>-0.536461</td>\n",
       "      <td>-1.177414</td>\n",
       "      <td>-0.827452</td>\n",
       "      <td>0.261866</td>\n",
       "      <td>-1.390910</td>\n",
       "      <td>1.331729</td>\n",
       "      <td>-1.626941</td>\n",
       "      <td>-0.077886</td>\n",
       "      <td>3.164828</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>-0.282896</td>\n",
       "      <td>-1.149107</td>\n",
       "      <td>-0.827452</td>\n",
       "      <td>0.244511</td>\n",
       "      <td>-1.418035</td>\n",
       "      <td>1.429997</td>\n",
       "      <td>-1.716847</td>\n",
       "      <td>-0.077886</td>\n",
       "      <td>1.148723</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>-0.240635</td>\n",
       "      <td>-1.120799</td>\n",
       "      <td>-0.827452</td>\n",
       "      <td>0.244511</td>\n",
       "      <td>-1.472284</td>\n",
       "      <td>1.498785</td>\n",
       "      <td>-1.724339</td>\n",
       "      <td>-0.077886</td>\n",
       "      <td>0.261637</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11917</th>\n",
       "      <td>0.393278</td>\n",
       "      <td>1.257029</td>\n",
       "      <td>1.802581</td>\n",
       "      <td>-0.848836</td>\n",
       "      <td>-1.418035</td>\n",
       "      <td>1.606881</td>\n",
       "      <td>-0.907690</td>\n",
       "      <td>-0.077886</td>\n",
       "      <td>-0.544805</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11917</th>\n",
       "      <td>0.689104</td>\n",
       "      <td>1.341951</td>\n",
       "      <td>2.153252</td>\n",
       "      <td>-0.848836</td>\n",
       "      <td>-1.517492</td>\n",
       "      <td>1.597054</td>\n",
       "      <td>-0.892705</td>\n",
       "      <td>-0.077886</td>\n",
       "      <td>-0.302872</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11917</th>\n",
       "      <td>0.562321</td>\n",
       "      <td>1.228721</td>\n",
       "      <td>1.627246</td>\n",
       "      <td>-0.848836</td>\n",
       "      <td>-1.535576</td>\n",
       "      <td>1.557746</td>\n",
       "      <td>-0.862736</td>\n",
       "      <td>-0.077886</td>\n",
       "      <td>-0.625449</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11917</th>\n",
       "      <td>0.942669</td>\n",
       "      <td>1.341951</td>\n",
       "      <td>2.591591</td>\n",
       "      <td>-0.779417</td>\n",
       "      <td>-1.607908</td>\n",
       "      <td>1.528266</td>\n",
       "      <td>-0.832768</td>\n",
       "      <td>-0.077886</td>\n",
       "      <td>-1.028670</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11917</th>\n",
       "      <td>1.069451</td>\n",
       "      <td>1.483489</td>\n",
       "      <td>3.117598</td>\n",
       "      <td>-0.727353</td>\n",
       "      <td>-1.598867</td>\n",
       "      <td>1.498785</td>\n",
       "      <td>-0.825275</td>\n",
       "      <td>-0.077886</td>\n",
       "      <td>0.100349</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>286032 rows × 9 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "          dim_0     dim_1     dim_2     dim_3     dim_4     dim_5     dim_6  \\\n",
       "0     -0.578722 -1.262337 -0.827452  0.331285 -1.354744  1.282595 -1.671894   \n",
       "0     -0.578722 -1.262337 -0.827452  0.331285 -1.390910  1.302248 -1.626941   \n",
       "0     -0.536461 -1.177414 -0.827452  0.261866 -1.390910  1.331729 -1.626941   \n",
       "0     -0.282896 -1.149107 -0.827452  0.244511 -1.418035  1.429997 -1.716847   \n",
       "0     -0.240635 -1.120799 -0.827452  0.244511 -1.472284  1.498785 -1.724339   \n",
       "...         ...       ...       ...       ...       ...       ...       ...   \n",
       "11917  0.393278  1.257029  1.802581 -0.848836 -1.418035  1.606881 -0.907690   \n",
       "11917  0.689104  1.341951  2.153252 -0.848836 -1.517492  1.597054 -0.892705   \n",
       "11917  0.562321  1.228721  1.627246 -0.848836 -1.535576  1.557746 -0.862736   \n",
       "11917  0.942669  1.341951  2.591591 -0.779417 -1.607908  1.528266 -0.832768   \n",
       "11917  1.069451  1.483489  3.117598 -0.727353 -1.598867  1.498785 -0.825275   \n",
       "\n",
       "          dim_7     dim_8  \n",
       "0     -0.077886  2.197098  \n",
       "0     -0.077886  2.439030  \n",
       "0     -0.077886  3.164828  \n",
       "0     -0.077886  1.148723  \n",
       "0     -0.077886  0.261637  \n",
       "...         ...       ...  \n",
       "11917 -0.077886 -0.544805  \n",
       "11917 -0.077886 -0.302872  \n",
       "11917 -0.077886 -0.625449  \n",
       "11917 -0.077886 -1.028670  \n",
       "11917 -0.077886  0.100349  \n",
       "\n",
       "[286032 rows x 9 columns]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.feature_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_len = 24\n",
    "out_size = 1\n",
    "out_dim = 1\n",
    "# config['data_window_len'] = max_len\n",
    "# config['task'] = 'simmtm'\n",
    "# config['normalization_layer'] = 'BatchNorm'\n",
    "# config['out_len'] = 24\n",
    "# config['out_dim'] = 7\n",
    "# config['d_model'] = 16\n",
    "# config['dim_feedforward'] = 128\n",
    "# config['num_heads'] = 4\n",
    "# config['num_layers'] = 1\n",
    "# from models.ts_transformer import model_factory\n",
    "# model = model_factory(config, data)\n",
    "model = DemoSimMTMTransformerEncoder(max_len=max_len, feat_dim=data.feature_df.shape[1], out_len=out_size, out_dim=out_dim, \n",
    "                                     d_model=4, n_heads=4, num_layers=2, dim_feedforward=8)\n",
    "\n",
    "device = \"cuda\"\n",
    "model.to(device)\n",
    "model.tau = 0.05\n",
    "model.mask_length = max_len//2\n",
    "model.mask_ratio = 0.5\n",
    "model.temporal_unit = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "batch_size = 64\n",
    "\n",
    "train_indices, val_indices, _ = split_dataset(data_indices=data.all_IDs,\n",
    "                                                         validation_method='ShuffleSplit',\n",
    "                                                         n_splits=1,\n",
    "                                                         validation_ratio=0.2,\n",
    "                                                         test_set_ratio=0,  # used only if test_indices not explicitly specified\n",
    "                                                         test_indices=None,\n",
    "                                                         random_seed=1337,\n",
    "                                                         labels=None)\n",
    "train_indices = train_indices[0]\n",
    "val_indices = val_indices[0]\n",
    "test_indices = np.array(test_data.all_IDs)\n",
    "\n",
    "train_dataloader = DataLoader(train_indices, batch_size=batch_size, shuffle=True)\n",
    "val_dataloader = DataLoader(val_indices, batch_size=batch_size, shuffle=True)\n",
    "test_dataloader = DataLoader(test_indices, batch_size=batch_size, shuffle=True)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Self-Supervised Learning Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 0 - Training loss: 4.23 - MSE loss: 0.74 - Contrastive loss: 4.67: 100%|██████████████████████████████████████████████████████████████████████| 149/149 [00:10<00:00, 13.99it/s]\n",
      "Epoch 1 - Training loss: 4.11 - MSE loss: 0.71 - Contrastive loss: 4.54:   1%|▉                                                                       | 2/149 [00:00<00:10, 14.02it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0 - Training loss: 6.78 0.98 6.53 - Validation loss: 4.10 0.82 4.41\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 1 - Training loss: 3.20 - MSE loss: 0.62 - Contrastive loss: 4.07: 100%|██████████████████████████████████████████████████████████████████████| 149/149 [00:10<00:00, 13.98it/s]\n",
      "Epoch 2 - Training loss: 3.24 - MSE loss: 0.66 - Contrastive loss: 4.11:   1%|▉                                                                       | 2/149 [00:00<00:10, 14.37it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 - Training loss: 3.73 0.81 4.34 - Validation loss: 3.17 0.74 3.88\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 2 - Training loss: 2.93 - MSE loss: 0.65 - Contrastive loss: 3.98: 100%|██████████████████████████████████████████████████████████████████████| 149/149 [00:10<00:00, 14.11it/s]\n",
      "Epoch 3 - Training loss: 3.14 - MSE loss: 1.00 - Contrastive loss: 3.94:   1%|▉                                                                       | 2/149 [00:00<00:10, 14.33it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2 - Training loss: 3.13 0.77 4.00 - Validation loss: 2.85 0.71 3.75\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 3 - Training loss: 2.64 - MSE loss: 0.57 - Contrastive loss: 3.82: 100%|██████████████████████████████████████████████████████████████████████| 149/149 [00:10<00:00, 14.01it/s]\n",
      "Epoch 4 - Training loss: 2.76 - MSE loss: 0.70 - Contrastive loss: 3.90:   1%|▉                                                                       | 2/149 [00:00<00:10, 14.36it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3 - Training loss: 2.86 0.74 3.90 - Validation loss: 2.66 0.69 3.68\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 4 - Training loss: 2.53 - MSE loss: 0.61 - Contrastive loss: 3.73: 100%|██████████████████████████████████████████████████████████████████████| 149/149 [00:10<00:00, 14.11it/s]\n",
      "Epoch 5 - Training loss: 2.65 - MSE loss: 0.76 - Contrastive loss: 3.78:   1%|▉                                                                       | 2/149 [00:00<00:10, 14.32it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4 - Training loss: 2.70 0.73 3.84 - Validation loss: 2.55 0.68 3.66\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 5 - Training loss: 2.55 - MSE loss: 0.79 - Contrastive loss: 3.65: 100%|██████████████████████████████████████████████████████████████████████| 149/149 [00:10<00:00, 14.19it/s]\n",
      "Epoch 6 - Training loss: 2.49 - MSE loss: 0.64 - Contrastive loss: 3.77:   1%|▉                                                                       | 2/149 [00:00<00:10, 13.99it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 5 - Training loss: 2.59 0.72 3.80 - Validation loss: 2.45 0.67 3.63\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 6 - Training loss: 2.41 - MSE loss: 0.66 - Contrastive loss: 3.68: 100%|██████████████████████████████████████████████████████████████████████| 149/149 [00:10<00:00, 14.00it/s]\n",
      "Epoch 7 - Training loss: 2.44 - MSE loss: 0.65 - Contrastive loss: 3.80:   1%|▉                                                                       | 2/149 [00:00<00:10, 14.39it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 6 - Training loss: 2.50 0.71 3.78 - Validation loss: 2.39 0.66 3.62\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 7 - Training loss: 2.63 - MSE loss: 1.08 - Contrastive loss: 3.60: 100%|██████████████████████████████████████████████████████████████████████| 149/149 [00:10<00:00, 13.99it/s]\n",
      "Epoch 8 - Training loss: 2.33 - MSE loss: 0.57 - Contrastive loss: 3.78:   1%|▉                                                                       | 2/149 [00:00<00:10, 14.12it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 7 - Training loss: 2.44 0.70 3.76 - Validation loss: 2.33 0.65 3.59\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 8 - Training loss: 2.40 - MSE loss: 0.79 - Contrastive loss: 3.64: 100%|██████████████████████████████████████████████████████████████████████| 149/149 [00:10<00:00, 13.91it/s]\n",
      "Epoch 9 - Training loss: 2.32 - MSE loss: 0.65 - Contrastive loss: 3.69:   1%|▉                                                                       | 2/149 [00:00<00:10, 14.39it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 8 - Training loss: 2.39 0.70 3.75 - Validation loss: 2.29 0.65 3.59\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 9 - Training loss: 2.22 - MSE loss: 0.55 - Contrastive loss: 3.70: 100%|██████████████████████████████████████████████████████████████████████| 149/149 [00:10<00:00, 13.88it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 9 - Training loss: 2.35 0.70 3.73 - Validation loss: 2.27 0.65 3.60\n",
      "Best Epoch 9 - Best Validation loss: 2.2673282623291016\n"
     ]
    }
   ],
   "source": [
    "i = 0\n",
    "max_epoch = 10\n",
    "best_loss = 1e10\n",
    "best_epoch = 0\n",
    "device = \"cuda\"\n",
    "loss_fn = nn.MSELoss()\n",
    "best_model = copy.deepcopy(model)\n",
    "\n",
    "while i < max_epoch:\n",
    "    train_loss = { \"loss\": [], \"loss_mse\": [], \"loss_con\": []}\n",
    "    progress_bar = tqdm(train_dataloader)\n",
    "    \n",
    "    for IDs in progress_bar:\n",
    "        model.train()\n",
    "        X = torch.tensor(data.feature_df.loc[IDs].to_numpy()).to(device)\n",
    "        X = X.float()\n",
    "        X = X.reshape(-1, max_len, X.shape[-1])\n",
    "        # X = X[:, :, -1:]\n",
    "        \n",
    "        pred, s = model(X)  # (batch_size, padded_length, feat_dim)\n",
    "        \n",
    "        loss_mse = loss_fn(pred, X) \n",
    "\n",
    "        loss_con = demo_contrastive_loss(s, X.shape[0])\n",
    "\n",
    "        loss = 1/(model.w1.pow(2)) * loss_mse + 1/(model.w2.pow(2)) * loss_con + torch.log(model.w1) + torch.log(model.w2)\n",
    "  \n",
    "\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "\n",
    "        nn.utils.clip_grad_norm_(model.parameters(), max_norm=4.0)\n",
    "        optimizer.step()\n",
    "        # import ipdb; ipdb.set_trace()\n",
    "        progress_bar.set_description(\"Epoch {0} - Training loss: {1:.2f} - MSE loss: {2:.2f} - Contrastive loss: {3:.2f}\".format(i, \n",
    "                loss.cpu().detach().numpy().item(), loss_mse.cpu().detach().numpy().item(), loss_con.cpu().detach().numpy().item())) \n",
    "        train_loss[\"loss\"].append(loss)\n",
    "        train_loss[\"loss_mse\"].append(loss_mse)\n",
    "        train_loss[\"loss_con\"].append(loss_con)\n",
    "    \n",
    "            \n",
    "    with torch.no_grad():\n",
    "        val_loss = { \"loss\": [], \"loss_mse\": [], \"loss_con\": []}\n",
    "        for IDs in val_dataloader:\n",
    "            model.eval()\n",
    "            X = torch.tensor(data.feature_df.loc[IDs].to_numpy()).to(device)\n",
    "            X = X.float()\n",
    "            X = X.reshape(-1, max_len, X.shape[-1])\n",
    "            # X = X[:, :, -1:]\n",
    "\n",
    "\n",
    "            pred, s = model(X)  # (batch_size, padded_length, feat_dim)\n",
    "        \n",
    "            loss_mse = loss_fn(pred, X) \n",
    "\n",
    "            loss_con = demo_contrastive_loss(s, X.shape[0])\n",
    "\n",
    "            loss = 1/(model.w1.pow(2)) * loss_mse + 1/(model.w2.pow(2)) * loss_con + torch.log(model.w1) + torch.log(model.w2)\n",
    "\n",
    "            val_loss[\"loss\"].append(loss)\n",
    "            val_loss[\"loss_mse\"].append(loss_mse)\n",
    "            val_loss[\"loss_con\"].append(loss_con)\n",
    "\n",
    "        train_loss[\"loss\"] = torch.tensor(train_loss[\"loss\"]).mean()\n",
    "        train_loss[\"loss_mse\"] = torch.tensor(train_loss[\"loss_mse\"]).mean()\n",
    "        train_loss[\"loss_con\"] = torch.tensor(train_loss[\"loss_con\"]).mean()\n",
    "        val_loss[\"loss\"] = torch.tensor(val_loss[\"loss\"]).mean()\n",
    "        val_loss[\"loss_mse\"] = torch.tensor(val_loss[\"loss_mse\"]).mean()\n",
    "        val_loss[\"loss_con\"] = torch.tensor(val_loss[\"loss_con\"]).mean()\n",
    "\n",
    "        if val_loss[\"loss\"] < best_loss:\n",
    "            best_loss = val_loss[\"loss\"]\n",
    "            best_model = copy.deepcopy(model)\n",
    "            best_epoch = i\n",
    "    \n",
    "        progress_bar.write(\"Epoch {0} - Training loss: {1:.2f} {2:.2f} {3:.2f} - Validation loss: {4:.2f} {5:.2f} {6:.2f}\".format(i, \n",
    "            train_loss[\"loss\"].cpu().detach().numpy().item(), train_loss[\"loss_mse\"].cpu().detach().numpy().item(), train_loss[\"loss_con\"].cpu().detach().numpy().item(),\n",
    "            val_loss[\"loss\"].cpu().detach().numpy().item(), val_loss[\"loss_mse\"].cpu().detach().numpy().item(), val_loss[\"loss_con\"].cpu().detach().numpy().item()))\n",
    "    i += 1\n",
    "    \n",
    "    \n",
    "tqdm.write(\"Best Epoch {} - Best Validation loss: {}\".format(best_epoch, best_loss))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Finetune Training Loop\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "finetune_model = copy.deepcopy(best_model)\n",
    "optimizer = torch.optim.AdamW(finetune_model.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 0 - Training loss: 34837.46: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████| 149/149 [00:07<00:00, 20.94it/s]\n",
      "Epoch 1 - Training loss: 27544.54:   1%|█▍                                                                                                            | 2/149 [00:00<00:07, 19.43it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0 - Training loss: 31422.14 - Validation loss: 29367.37\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 1 - Training loss: 22497.54: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████| 149/149 [00:07<00:00, 20.65it/s]\n",
      "Epoch 2 - Training loss: 21008.80:   1%|█▍                                                                                                            | 2/149 [00:00<00:07, 19.49it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 - Training loss: 28681.35 - Validation loss: 25832.41\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 2 - Training loss: 16309.61: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████| 149/149 [00:07<00:00, 20.16it/s]\n",
      "Epoch 3 - Training loss: 15380.78:   1%|█▍                                                                                                            | 2/149 [00:00<00:08, 17.73it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2 - Training loss: 23916.93 - Validation loss: 20252.27\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 3 - Training loss: 14280.37: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████| 149/149 [00:07<00:00, 20.03it/s]\n",
      "Epoch 4 - Training loss: 16303.24:   2%|██▏                                                                                                           | 3/149 [00:00<00:06, 20.97it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3 - Training loss: 18323.24 - Validation loss: 15602.92\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 4 - Training loss: 16054.93: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████| 149/149 [00:07<00:00, 20.53it/s]\n",
      "Epoch 5 - Training loss: 10605.63:   2%|██▏                                                                                                           | 3/149 [00:00<00:07, 20.55it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4 - Training loss: 13232.80 - Validation loss: 9963.98\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 5 - Training loss: 7481.14: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 149/149 [00:07<00:00, 21.17it/s]\n",
      "Epoch 6 - Training loss: 5631.07:   2%|██▏                                                                                                            | 3/149 [00:00<00:06, 21.24it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 5 - Training loss: 9234.78 - Validation loss: 6904.74\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 6 - Training loss: 15172.48: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████| 149/149 [00:07<00:00, 21.23it/s]\n",
      "Epoch 7 - Training loss: 12741.97:   2%|██▏                                                                                                           | 3/149 [00:00<00:06, 21.29it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 6 - Training loss: 7900.12 - Validation loss: 5592.28\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 7 - Training loss: 4540.79: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 149/149 [00:07<00:00, 21.25it/s]\n",
      "Epoch 8 - Training loss: 7992.80:   2%|██▏                                                                                                            | 3/149 [00:00<00:06, 21.28it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 7 - Training loss: 7426.88 - Validation loss: 4517.94\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 8 - Training loss: 5199.29: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 149/149 [00:06<00:00, 21.31it/s]\n",
      "Epoch 9 - Training loss: 6122.55:   2%|██▏                                                                                                            | 3/149 [00:00<00:06, 21.24it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 8 - Training loss: 7415.72 - Validation loss: 4677.93\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 9 - Training loss: 7548.82: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 149/149 [00:07<00:00, 21.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 9 - Training loss: 7343.61 - Validation loss: 4388.68\n",
      "Best Epoch 9 - Best Validation loss: 4388.6767578125\n"
     ]
    }
   ],
   "source": [
    "i = 0\n",
    "max_epoch = 10\n",
    "best_loss = 1e10\n",
    "best_finetune_model = copy.deepcopy(best_model)\n",
    "best_epoch = 0\n",
    "device = \"cuda\"\n",
    "finetune_model.to(device)\n",
    "while i < max_epoch:\n",
    "    train_loss = []\n",
    "    progress_bar = tqdm(train_dataloader)\n",
    "    \n",
    "    for IDs in progress_bar:\n",
    "        finetune_model.train()\n",
    "        \n",
    "        X = torch.tensor(data.feature_df.loc[IDs].to_numpy()).to(device)\n",
    "        X = X.reshape(-1, max_len, X.shape[-1])\n",
    "        targets = torch.tensor(data.labels_df.loc[IDs].to_numpy()).to(device)\n",
    "        targets = targets.reshape(targets.shape[0], out_size, -1)\n",
    "        \n",
    "        pred = finetune_model.predict(X.float())\n",
    "        pred = pred.reshape(X.shape[0], out_size, -1)\n",
    "        loss = loss_fn(pred, targets)\n",
    "\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "\n",
    "        nn.utils.clip_grad_norm_(finetune_model.parameters(), max_norm=4.0)\n",
    "        optimizer.step()\n",
    "\n",
    "        progress_bar.set_description(\"Epoch {} - Training loss: {:.2f}\".format(i, loss)) \n",
    "        train_loss.append(loss)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        val_loss = []\n",
    "        for IDs in val_dataloader:\n",
    "            finetune_model.eval()\n",
    "            X = torch.tensor(data.feature_df.loc[IDs].to_numpy()).to(device)\n",
    "            X = X.reshape(-1, max_len, X.shape[-1])\n",
    "            targets = torch.tensor(data.labels_df.loc[IDs].to_numpy()).to(device)\n",
    "            targets = targets.reshape(targets.shape[0], out_size, -1)\n",
    "\n",
    "            pred = finetune_model.predict(X.float())\n",
    "            pred = pred.reshape(X.shape[0], out_size, -1)\n",
    "            \n",
    "            loss = loss_fn(pred, targets)\n",
    "            val_loss.append(loss)\n",
    "\n",
    "        train_loss = torch.tensor(train_loss).mean()\n",
    "        val_loss = torch.tensor(val_loss).mean()\n",
    "\n",
    "        if val_loss < best_loss:\n",
    "            best_loss = val_loss\n",
    "            best_finetune_model = copy.deepcopy(finetune_model)\n",
    "            best_epoch = i\n",
    "    \n",
    "    progress_bar.write(\"Epoch {} - Training loss: {:.2f} - Validation loss: {:.2f}\".format(i, train_loss, val_loss))\n",
    "    i += 1\n",
    "    \n",
    "    \n",
    "tqdm.write(\"Best Epoch {} - Best Validation loss: {}\".format(best_epoch, best_loss))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test loss: 4371.2177734375\n"
     ]
    }
   ],
   "source": [
    "test_loss = []\n",
    "with torch.no_grad():\n",
    "    for IDs in test_dataloader:\n",
    "        best_finetune_model.eval()\n",
    "        X = torch.tensor(data.feature_df.loc[IDs].to_numpy()).to(device)\n",
    "        X = X.reshape(-1, max_len, X.shape[-1])\n",
    "        targets = torch.tensor(data.labels_df.loc[IDs].to_numpy()).to(device)\n",
    "        targets = targets.reshape(targets.shape[0], out_size, -1)\n",
    "        \n",
    "\n",
    "        pred = best_finetune_model.predict(X.float())\n",
    "        pred = pred.reshape(X.shape[0], out_size, -1)\n",
    "        loss = loss_fn(pred, targets)\n",
    "\n",
    "\n",
    "        test_loss.append(loss)\n",
    "\n",
    "\n",
    "test_loss = torch.tensor(test_loss).mean()\n",
    "print(\"Test MSE loss: {}\".format(test_loss))\n",
    "print(\"Test RMSE loss: {}\".format(np.sqrt(test_loss)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "<center><img src=\"../img/img_8.PNG\"/><center/>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reference:\n",
    "1. https://arxiv.org/abs/2302.00861\n",
    "2. https://github.com/gzerveas/mvts_transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "No Pretrain\n",
    "\n",
    "Test loss: 4083.34765625\n",
    "\n",
    "Test loss: 4096.765625\n",
    "\n",
    "Test loss: 3880.41064453125\n",
    "\n",
    "Test loss: 3818.454833984375\n",
    "\n",
    "Test loss: 3640.204833984375\n",
    "\n",
    "Test loss: 3793.255615234375\n",
    "\n",
    "Test loss: 3386.54638671875\n",
    "\n",
    "Test loss: 3532.69140625\n",
    "\n",
    "Test loss: 3568.943115234375\n",
    "\n",
    "Test loss: 3639.838623046875\n",
    "\n",
    "Pretrain\n",
    "\n",
    "Test loss: 3398.438720703125\n",
    "\n",
    "Test loss: 3389.167724609375\n",
    "\n",
    "Test loss: 3470.773681640625\n",
    "\n",
    "Test loss: 3443.25830078125\n",
    "\n",
    "Test loss: 3331.67919921875\n",
    "\n",
    "Test loss: 3753.1474609375\n",
    "\n",
    "Test loss: 3480.992431640625\n",
    "\n",
    "Test loss: 3532.7275390625\n",
    "\n",
    "Test loss: 3068.515869140625\n",
    "\n",
    "Test loss: 3118.0390625"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}