lstm_baseline.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch_geometric\n",
    "import networkx as nx\n",
    "import os\n",
    "import torch\n",
    "from torch_geometric.loader import DataLoader\n",
    "from tqdm import tqdm\n",
    "from torch_geometric.data import Data\n",
    "import pandas as pd\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#>>> pd.read_excel('data/RhythmNames.xlsx')\n",
    "#   Acronym Name                                     Full Name\n",
    "#0          SB                               Sinus Bradycardia\n",
    "#1          SR                                    Sinus Rhythm\n",
    "#2         AFIB                            Atrial Fibrillation\n",
    "#3         ST                                Sinus Tachycardia\n",
    "#4          AF                                  Atrial Flutter\n",
    "#5          SI                              Sinus Irregularity\n",
    "#6          SVT                   Supraventricular Tachycardia\n",
    "#7          AT                              Atrial Tachycardia\n",
    "#8         AVNRT  Atrioventricular  Node Reentrant Tachycardia\n",
    "#9          AVRT        Atrioventricular Reentrant Tachycardia\n",
    "#10        SAAWR       Sinus Atrium to Atrial Wandering Rhythm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#>>> pd.read_excel('data/ConditionNames.xlsx')\n",
    "#   Acronym Name                                  Full Name\n",
    "#0          1AVB            1 degree atrioventricular block\n",
    "#1          2AVB            2 degree atrioventricular block\n",
    "#2         2AVB1  2 degree atrioventricular block(Type one)\n",
    "#3         2AVB2  2 degree atrioventricular block(Type two)\n",
    "#4          3AVB            3 degree atrioventricular block\n",
    "#5           ABI                            atrial bigeminy\n",
    "#6           ALS                            Axis left shift\n",
    "#7           APB                     atrial premature beats\n",
    "#8           AQW                            abnormal Q wave\n",
    "#9           ARS                           Axis right shift\n",
    "#10          AVB                     atrioventricular block\n",
    "#11          CCR                 countercolockwise rotation\n",
    "#12           CR                        colockwise rotation\n",
    "#13          ERV     Early repolarization of the ventricles\n",
    "#14         FQRS                                  fQRS Wave\n",
    "#15          IDC            Interior differences conduction\n",
    "#16          IVB                     Intraventricular block\n",
    "#17          JEB                     junctional escape beat\n",
    "#18          JPS                              J point shift\n",
    "#19          JPT                  junctional premature beat\n",
    "#20         LBBB                   left bundle branch block\n",
    "#21        LBBBB              left back bundle branch block\n",
    "#22        LFBBB             left front bundle branch block\n",
    "#23         LRRI                           Long RR interval\n",
    "#24          LVH                 left ventricle hypertrophy\n",
    "#25         LVHV                left ventricle high voltage\n",
    "#26      LVQRSAL              lower voltage QRS in all lead\n",
    "#27      LVQRSCL            lower voltage QRS in chest lead\n",
    "#28      LVQRSLL             lower voltage QRS in limb lead\n",
    "#29           MI                      myocardial infarction\n",
    "#30         MIBW         myocardial infraction in back wall\n",
    "#31         MIFW   Myocardial infgraction in the front wall\n",
    "#32         MILW    Myocardial infraction in the lower wall\n",
    "#33         MISW     Myocardial infraction in the side wall\n",
    "#34         PRIE                      PR interval extension\n",
    "#35          PWC                              P wave Change\n",
    "#36         QTIE                      QT interval extension\n",
    "#37          RAH                   right atrial hypertrophy\n",
    "#38         RAHV                  right atrial high voltage\n",
    "#39         RBBB                  right bundle branch block\n",
    "#40          RVH                right ventricle hypertrophy\n",
    "#41         STDD                               ST drop down\n",
    "#42          STE                               ST extension\n",
    "#43         STTC                                ST-T Change\n",
    "#44         STTU                                 ST tilt up\n",
    "#45          TWC                              T wave Change\n",
    "#46          TWO                            T wave opposite\n",
    "#47           UW                                     U wave\n",
    "#48           VB                       ventricular bigeminy\n",
    "#49          VEB                    ventricular escape beat\n",
    "#50          VFW                    ventricular fusion wave\n",
    "#51          VPB                 ventricular premature beat\n",
    "#52          VPE                  ventricular preexcitation\n",
    "#53          VET               ventricular escape trigeminy\n",
    "#54         WAVN   Wandering in the atrioventricalualr node\n",
    "#55          WPW                                        WPW"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "diagnostics_df = pd.read_excel('data/Diagnostics.xlsx')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import LabelEncoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LabelEncoder()"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "le = LabelEncoder()\n",
    "le.fit(diagnostics_df.Rhythm.unique())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['AF', 'AFIB', 'AT', 'AVNRT', 'AVRT', 'SA', 'SAAWR', 'SB', 'SR',\n",
       "       'ST', 'SVT'], dtype=object)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "le.classes_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "diagnostics_df['label'] = le.transform(diagnostics_df.Rhythm.values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0         1\n",
       "1         7\n",
       "2         5\n",
       "3         7\n",
       "4         0\n",
       "         ..\n",
       "10641    10\n",
       "10642    10\n",
       "10643    10\n",
       "10644    10\n",
       "10645    10\n",
       "Name: label, Length: 10646, dtype: int64"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "diagnostics_df['label']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 53%|█████▎    | 5690/10653 [01:01<00:50, 99.12it/s] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MUSE_20180113_124215_52000.csv\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10653/10653 [01:52<00:00, 94.39it/s]\n"
     ]
    }
   ],
   "source": [
    "data_list = []\n",
    "for file_name in tqdm(os.listdir('data/')):\n",
    "    if file_name.endswith('.csv'):\n",
    "        sub_name = file_name.split('.')[0]\n",
    "        y = torch.tensor(diagnostics_df.loc[diagnostics_df.FileName == sub_name].label.values[0])\n",
    "        df = pd.read_csv(f'data/{file_name}', header=None)\n",
    "        if len(df) != 5000:\n",
    "            print(file_name)\n",
    "            continue\n",
    "        x = torch.Tensor(np.array([df[col].values for col in df.columns])).t()\n",
    "        data = [x, y]\n",
    "        data_list.append(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_epochs = 250\n",
    "lr = 0.01\n",
    "n_folds = 5\n",
    "lstm_input_size = 12\n",
    "hidden_state_size = 30\n",
    "batch_size = 32\n",
    "num_sequence_layers = 2\n",
    "output_dim = 11\n",
    "num_time_steps = 5000\n",
    "rnn_type = 'LSTM'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = DataLoader(data_list[:8000], batch_size=batch_size, shuffle=True)\n",
    "test_loader = DataLoader(data_list[8000:], batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "instance = next(iter(train_loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 5000, 12])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "instance[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "#g = torch_geometric.utils.to_networkx(instance, to_undirected=True)\n",
    "#nx.draw(g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 1:\n",
      "=======\n",
      "Batch 0: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 2:\n",
      "=======\n",
      "Batch 1: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 3:\n",
      "=======\n",
      "Batch 2: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 4:\n",
      "=======\n",
      "Batch 3: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 5:\n",
      "=======\n",
      "Batch 4: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 6:\n",
      "=======\n",
      "Batch 5: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 7:\n",
      "=======\n",
      "Batch 6: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 8:\n",
      "=======\n",
      "Batch 7: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 9:\n",
      "=======\n",
      "Batch 8: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 10:\n",
      "=======\n",
      "Batch 9: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 11:\n",
      "=======\n",
      "Batch 10: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 12:\n",
      "=======\n",
      "Batch 11: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 13:\n",
      "=======\n",
      "Batch 12: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 14:\n",
      "=======\n",
      "Batch 13: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 15:\n",
      "=======\n",
      "Batch 14: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 16:\n",
      "=======\n",
      "Batch 15: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 17:\n",
      "=======\n",
      "Batch 16: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 18:\n",
      "=======\n",
      "Batch 17: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 19:\n",
      "=======\n",
      "Batch 18: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 20:\n",
      "=======\n",
      "Batch 19: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 21:\n",
      "=======\n",
      "Batch 20: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 22:\n",
      "=======\n",
      "Batch 21: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 23:\n",
      "=======\n",
      "Batch 22: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 24:\n",
      "=======\n",
      "Batch 23: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 25:\n",
      "=======\n",
      "Batch 24: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 26:\n",
      "=======\n",
      "Batch 25: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 27:\n",
      "=======\n",
      "Batch 26: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 28:\n",
      "=======\n",
      "Batch 27: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 29:\n",
      "=======\n",
      "Batch 28: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 30:\n",
      "=======\n",
      "Batch 29: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 31:\n",
      "=======\n",
      "Batch 30: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 32:\n",
      "=======\n",
      "Batch 31: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 33:\n",
      "=======\n",
      "Batch 32: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 34:\n",
      "=======\n",
      "Batch 33: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 35:\n",
      "=======\n",
      "Batch 34: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 36:\n",
      "=======\n",
      "Batch 35: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 37:\n",
      "=======\n",
      "Batch 36: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 38:\n",
      "=======\n",
      "Batch 37: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 39:\n",
      "=======\n",
      "Batch 38: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 40:\n",
      "=======\n",
      "Batch 39: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 41:\n",
      "=======\n",
      "Batch 40: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 42:\n",
      "=======\n",
      "Batch 41: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 43:\n",
      "=======\n",
      "Batch 42: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 44:\n",
      "=======\n",
      "Batch 43: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 45:\n",
      "=======\n",
      "Batch 44: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 46:\n",
      "=======\n",
      "Batch 45: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 47:\n",
      "=======\n",
      "Batch 46: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 48:\n",
      "=======\n",
      "Batch 47: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 49:\n",
      "=======\n",
      "Batch 48: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 50:\n",
      "=======\n",
      "Batch 49: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 51:\n",
      "=======\n",
      "Batch 50: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 52:\n",
      "=======\n",
      "Batch 51: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 53:\n",
      "=======\n",
      "Batch 52: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 54:\n",
      "=======\n",
      "Batch 53: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 55:\n",
      "=======\n",
      "Batch 54: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 56:\n",
      "=======\n",
      "Batch 55: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 57:\n",
      "=======\n",
      "Batch 56: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 58:\n",
      "=======\n",
      "Batch 57: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 59:\n",
      "=======\n",
      "Batch 58: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 60:\n",
      "=======\n",
      "Batch 59: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 61:\n",
      "=======\n",
      "Batch 60: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 62:\n",
      "=======\n",
      "Batch 61: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 63:\n",
      "=======\n",
      "Batch 62: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 64:\n",
      "=======\n",
      "Batch 63: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 65:\n",
      "=======\n",
      "Batch 64: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 66:\n",
      "=======\n",
      "Batch 65: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 67:\n",
      "=======\n",
      "Batch 66: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 68:\n",
      "=======\n",
      "Batch 67: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 69:\n",
      "=======\n",
      "Batch 68: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 70:\n",
      "=======\n",
      "Batch 69: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 71:\n",
      "=======\n",
      "Batch 70: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 72:\n",
      "=======\n",
      "Batch 71: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 73:\n",
      "=======\n",
      "Batch 72: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 74:\n",
      "=======\n",
      "Batch 73: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 75:\n",
      "=======\n",
      "Batch 74: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 76:\n",
      "=======\n",
      "Batch 75: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 77:\n",
      "=======\n",
      "Batch 76: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 78:\n",
      "=======\n",
      "Batch 77: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 79:\n",
      "=======\n",
      "Batch 78: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 80:\n",
      "=======\n",
      "Batch 79: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 81:\n",
      "=======\n",
      "Batch 80: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 82:\n",
      "=======\n",
      "Batch 81: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 83:\n",
      "=======\n",
      "Batch 82: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 84:\n",
      "=======\n",
      "Batch 83: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 85:\n",
      "=======\n",
      "Batch 84: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 86:\n",
      "=======\n",
      "Batch 85: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 87:\n",
      "=======\n",
      "Batch 86: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 88:\n",
      "=======\n",
      "Batch 87: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 89:\n",
      "=======\n",
      "Batch 88: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 90:\n",
      "=======\n",
      "Batch 89: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 91:\n",
      "=======\n",
      "Batch 90: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 92:\n",
      "=======\n",
      "Batch 91: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 93:\n",
      "=======\n",
      "Batch 92: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 94:\n",
      "=======\n",
      "Batch 93: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 95:\n",
      "=======\n",
      "Batch 94: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 96:\n",
      "=======\n",
      "Batch 95: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 97:\n",
      "=======\n",
      "Batch 96: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 98:\n",
      "=======\n",
      "Batch 97: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 99:\n",
      "=======\n",
      "Batch 98: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 100:\n",
      "=======\n",
      "Batch 99: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 101:\n",
      "=======\n",
      "Batch 100: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 102:\n",
      "=======\n",
      "Batch 101: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 103:\n",
      "=======\n",
      "Batch 102: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 104:\n",
      "=======\n",
      "Batch 103: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 105:\n",
      "=======\n",
      "Batch 104: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 106:\n",
      "=======\n",
      "Batch 105: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 107:\n",
      "=======\n",
      "Batch 106: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 108:\n",
      "=======\n",
      "Batch 107: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 109:\n",
      "=======\n",
      "Batch 108: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 110:\n",
      "=======\n",
      "Batch 109: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 111:\n",
      "=======\n",
      "Batch 110: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 112:\n",
      "=======\n",
      "Batch 111: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 113:\n",
      "=======\n",
      "Batch 112: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 114:\n",
      "=======\n",
      "Batch 113: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 115:\n",
      "=======\n",
      "Batch 114: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 116:\n",
      "=======\n",
      "Batch 115: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 117:\n",
      "=======\n",
      "Batch 116: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 118:\n",
      "=======\n",
      "Batch 117: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 119:\n",
      "=======\n",
      "Batch 118: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 120:\n",
      "=======\n",
      "Batch 119: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 121:\n",
      "=======\n",
      "Batch 120: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 122:\n",
      "=======\n",
      "Batch 121: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 123:\n",
      "=======\n",
      "Batch 122: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 124:\n",
      "=======\n",
      "Batch 123: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 125:\n",
      "=======\n",
      "Batch 124: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 126:\n",
      "=======\n",
      "Batch 125: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 127:\n",
      "=======\n",
      "Batch 126: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 128:\n",
      "=======\n",
      "Batch 127: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 129:\n",
      "=======\n",
      "Batch 128: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 130:\n",
      "=======\n",
      "Batch 129: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 131:\n",
      "=======\n",
      "Batch 130: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 132:\n",
      "=======\n",
      "Batch 131: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 133:\n",
      "=======\n",
      "Batch 132: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 134:\n",
      "=======\n",
      "Batch 133: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 135:\n",
      "=======\n",
      "Batch 134: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 136:\n",
      "=======\n",
      "Batch 135: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 137:\n",
      "=======\n",
      "Batch 136: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 138:\n",
      "=======\n",
      "Batch 137: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 139:\n",
      "=======\n",
      "Batch 138: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 140:\n",
      "=======\n",
      "Batch 139: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 141:\n",
      "=======\n",
      "Batch 140: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 142:\n",
      "=======\n",
      "Batch 141: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 143:\n",
      "=======\n",
      "Batch 142: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 144:\n",
      "=======\n",
      "Batch 143: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 145:\n",
      "=======\n",
      "Batch 144: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 146:\n",
      "=======\n",
      "Batch 145: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 147:\n",
      "=======\n",
      "Batch 146: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 148:\n",
      "=======\n",
      "Batch 147: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 149:\n",
      "=======\n",
      "Batch 148: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 150:\n",
      "=======\n",
      "Batch 149: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 151:\n",
      "=======\n",
      "Batch 150: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 152:\n",
      "=======\n",
      "Batch 151: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 153:\n",
      "=======\n",
      "Batch 152: torch.Size([32, 5000, 12])\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 154:\n",
      "=======\n",
      "Batch 153: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 155:\n",
      "=======\n",
      "Batch 154: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 156:\n",
      "=======\n",
      "Batch 155: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 157:\n",
      "=======\n",
      "Batch 156: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 158:\n",
      "=======\n",
      "Batch 157: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 159:\n",
      "=======\n",
      "Batch 158: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 160:\n",
      "=======\n",
      "Batch 159: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 161:\n",
      "=======\n",
      "Batch 160: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 162:\n",
      "=======\n",
      "Batch 161: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 163:\n",
      "=======\n",
      "Batch 162: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 164:\n",
      "=======\n",
      "Batch 163: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 165:\n",
      "=======\n",
      "Batch 164: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 166:\n",
      "=======\n",
      "Batch 165: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 167:\n",
      "=======\n",
      "Batch 166: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 168:\n",
      "=======\n",
      "Batch 167: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 169:\n",
      "=======\n",
      "Batch 168: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 170:\n",
      "=======\n",
      "Batch 169: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 171:\n",
      "=======\n",
      "Batch 170: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 172:\n",
      "=======\n",
      "Batch 171: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 173:\n",
      "=======\n",
      "Batch 172: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 174:\n",
      "=======\n",
      "Batch 173: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 175:\n",
      "=======\n",
      "Batch 174: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 176:\n",
      "=======\n",
      "Batch 175: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 177:\n",
      "=======\n",
      "Batch 176: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 178:\n",
      "=======\n",
      "Batch 177: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 179:\n",
      "=======\n",
      "Batch 178: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 180:\n",
      "=======\n",
      "Batch 179: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 181:\n",
      "=======\n",
      "Batch 180: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 182:\n",
      "=======\n",
      "Batch 181: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 183:\n",
      "=======\n",
      "Batch 182: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 184:\n",
      "=======\n",
      "Batch 183: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 185:\n",
      "=======\n",
      "Batch 184: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 186:\n",
      "=======\n",
      "Batch 185: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 187:\n",
      "=======\n",
      "Batch 186: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 188:\n",
      "=======\n",
      "Batch 187: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 189:\n",
      "=======\n",
      "Batch 188: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 190:\n",
      "=======\n",
      "Batch 189: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 191:\n",
      "=======\n",
      "Batch 190: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 192:\n",
      "=======\n",
      "Batch 191: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 193:\n",
      "=======\n",
      "Batch 192: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 194:\n",
      "=======\n",
      "Batch 193: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 195:\n",
      "=======\n",
      "Batch 194: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 196:\n",
      "=======\n",
      "Batch 195: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 197:\n",
      "=======\n",
      "Batch 196: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 198:\n",
      "=======\n",
      "Batch 197: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 199:\n",
      "=======\n",
      "Batch 198: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 200:\n",
      "=======\n",
      "Batch 199: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 201:\n",
      "=======\n",
      "Batch 200: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 202:\n",
      "=======\n",
      "Batch 201: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 203:\n",
      "=======\n",
      "Batch 202: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 204:\n",
      "=======\n",
      "Batch 203: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 205:\n",
      "=======\n",
      "Batch 204: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 206:\n",
      "=======\n",
      "Batch 205: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 207:\n",
      "=======\n",
      "Batch 206: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 208:\n",
      "=======\n",
      "Batch 207: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 209:\n",
      "=======\n",
      "Batch 208: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 210:\n",
      "=======\n",
      "Batch 209: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 211:\n",
      "=======\n",
      "Batch 210: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 212:\n",
      "=======\n",
      "Batch 211: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 213:\n",
      "=======\n",
      "Batch 212: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 214:\n",
      "=======\n",
      "Batch 213: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 215:\n",
      "=======\n",
      "Batch 214: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 216:\n",
      "=======\n",
      "Batch 215: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 217:\n",
      "=======\n",
      "Batch 216: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 218:\n",
      "=======\n",
      "Batch 217: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 219:\n",
      "=======\n",
      "Batch 218: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 220:\n",
      "=======\n",
      "Batch 219: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 221:\n",
      "=======\n",
      "Batch 220: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 222:\n",
      "=======\n",
      "Batch 221: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 223:\n",
      "=======\n",
      "Batch 222: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 224:\n",
      "=======\n",
      "Batch 223: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 225:\n",
      "=======\n",
      "Batch 224: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 226:\n",
      "=======\n",
      "Batch 225: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 227:\n",
      "=======\n",
      "Batch 226: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 228:\n",
      "=======\n",
      "Batch 227: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 229:\n",
      "=======\n",
      "Batch 228: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 230:\n",
      "=======\n",
      "Batch 229: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 231:\n",
      "=======\n",
      "Batch 230: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 232:\n",
      "=======\n",
      "Batch 231: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 233:\n",
      "=======\n",
      "Batch 232: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 234:\n",
      "=======\n",
      "Batch 233: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 235:\n",
      "=======\n",
      "Batch 234: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 236:\n",
      "=======\n",
      "Batch 235: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 237:\n",
      "=======\n",
      "Batch 236: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 238:\n",
      "=======\n",
      "Batch 237: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 239:\n",
      "=======\n",
      "Batch 238: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 240:\n",
      "=======\n",
      "Batch 239: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 241:\n",
      "=======\n",
      "Batch 240: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 242:\n",
      "=======\n",
      "Batch 241: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 243:\n",
      "=======\n",
      "Batch 242: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 244:\n",
      "=======\n",
      "Batch 243: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 245:\n",
      "=======\n",
      "Batch 244: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 246:\n",
      "=======\n",
      "Batch 245: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 247:\n",
      "=======\n",
      "Batch 246: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 248:\n",
      "=======\n",
      "Batch 247: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 249:\n",
      "=======\n",
      "Batch 248: torch.Size([32, 5000, 12])\n",
      "\n",
      "Step 250:\n",
      "=======\n",
      "Batch 249: torch.Size([32, 5000, 12])\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for step, data in enumerate(train_loader):\n",
    "    print(f'Step {step + 1}:')\n",
    "    print('=======')\n",
    "    print(f'Batch {step}: {data[0].shape}')\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Bi_RNN(nn.Module):\n",
    "\n",
    "    def __init__(self, input_dim, hidden_dim, batch_size, output_dim=11, num_layers=2, rnn_type='LSTM'):\n",
    "        super(Bi_RNN, self).__init__()\n",
    "        self.input_dim = input_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.batch_size = batch_size\n",
    "        self.num_layers = num_layers\n",
    "\n",
    "        #Define the initial linear hidden layer\n",
    "        self.init_linear = nn.Linear(self.input_dim, self.input_dim)\n",
    "\n",
    "        # Define the LSTM layer\n",
    "        self.lstm = eval('nn.' + rnn_type)(self.input_dim, self.hidden_dim, self.num_layers, batch_first=True, bidirectional=True)\n",
    "\n",
    "        # Define the output layer\n",
    "        self.linear = nn.Linear(self.hidden_dim * 2, output_dim)\n",
    "\n",
    "    def init_hidden(self):\n",
    "        # This is what we'll initialise our hidden state as\n",
    "        return (torch.zeros(self.num_layers, self.batch_size, self.hidden_dim),\n",
    "                torch.zeros(self.num_layers, self.batch_size, self.hidden_dim))\n",
    "\n",
    "    def forward(self, input):\n",
    "        #Forward pass through initial hidden layer\n",
    "        linear_input = self.init_linear(input)\n",
    "\n",
    "        # Forward pass through LSTM layer\n",
    "        # shape of lstm_out: [batch_size, input_size ,hidden_dim]\n",
    "        # shape of self.hidden: (a, b), where a and b both\n",
    "        # have shape (batch_size, num_layers, hidden_dim).\n",
    "        lstm_out, self.hidden = self.lstm(linear_input)\n",
    "        out_reduced = torch.cat((self.hidden[0][-2,:,:], self.hidden[0][-1,:,:]), dim = 1)\n",
    "        \n",
    "        \n",
    "        # Can pass on the entirety of lstm_out to the next layer if it is a seq2seq prediction\n",
    "        y_pred = self.linear(out_reduced)\n",
    "        # text_fea = torch.squeeze(text_fea, 1)\n",
    "        return y_pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    device = 'cuda:0'\n",
    "else:\n",
    "    device = 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "model1 = Bi_RNN(lstm_input_size, hidden_state_size, batch_size=batch_size, output_dim=output_dim, num_layers=num_sequence_layers, rnn_type=rnn_type).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bi_RNN(\n",
      "  (init_linear): Linear(in_features=12, out_features=12, bias=True)\n",
      "  (lstm): LSTM(12, 30, num_layers=2, batch_first=True, bidirectional=True)\n",
      "  (linear): Linear(in_features=60, out_features=11, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "print(model1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set model paramters and model type\n",
    "def set_model_parameters(model_type, lr=0.01):\n",
    "    model = model_type\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "    criterion = torch.nn.CrossEntropyLoss()\n",
    "    return model, optimizer, criterion\n",
    "\n",
    "# Train the model\n",
    "def train(model, optimizer,criterion):\n",
    "    model.train()\n",
    "\n",
    "    for idx, (x_batch, y_batch) in enumerate(train_loader):  # Iterate in batches over the training dataset.\n",
    "        x_batch = x_batch.view(-1, num_time_steps, lstm_input_size).to(device)\n",
    "        y_batch = y_batch.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        y_pred = model(x_batch)\n",
    "        loss = criterion(y_pred, y_batch)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()  # Clear gradients.\n",
    "\n",
    "# Test the model \n",
    "def test(loader, model):\n",
    "     model.eval()\n",
    "\n",
    "     correct = 0\n",
    "     for data in loader:  # Iterate in batches over the training/test dataset.\n",
    "         out = model(data[0].to(device))  \n",
    "         pred = out.argmax(dim=1)  # Use the class with highest probability.\n",
    "         correct += int((pred == data[1].to(device)).sum())  # Check against ground-truth labels.\n",
    "     return correct / len(loader.dataset)  # Derive ratio of correct predictions.\n",
    "\n",
    "# Training and Testing Pipeline \n",
    "def running_epochs(model,optimizer,criterion):\n",
    "    for epoch in range(12):\n",
    "        train(model,optimizer,criterion)\n",
    "        train_acc = test(train_loader, model)\n",
    "        test_acc = test(test_loader, model)\n",
    "        print(f'Epoch: {epoch:03d}, Train Acc: {train_acc:.4f}, Test Acc: {test_acc:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 000, Train Acc: 0.0428, Test Acc: 0.0389\n",
      "Epoch: 001, Train Acc: 0.0428, Test Acc: 0.0389\n"
     ]
    }
   ],
   "source": [
    "model, optimizer, criterion = set_model_parameters(model1, lr=0.001)\n",
    "running_epochs(model,optimizer,criterion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}