eeg_tensorflow.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# EEG Classification - Tensorflow\n",
    "updated: Sep. 01, 2018\n",
    "\n",
    "Data: https://www.physionet.org/pn4/eegmmidb/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Data Downloads\n",
    "\n",
    "### Warning: Executing these blocks will automatically create directories and download datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from keras import backend as K\n",
    "K.tensorflow_backend._get_available_gpus()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Tensorflow Style Guide\n",
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "# System\n",
    "import requests\n",
    "import re\n",
    "import os\n",
    "import pathlib\n",
    "import urllib\n",
    "\n",
    "# Modeling & Preprocessing\n",
    "from keras.layers import Conv2D, BatchNormalization, Activation, Flatten, Dense, Dropout, LSTM, Input, TimeDistributed\n",
    "from keras import initializers, Model, optimizers, callbacks\n",
    "from keras.utils.training_utils import multi_gpu_model\n",
    "from keras import backend as K\n",
    "from keras.models import load_model\n",
    "from keras.callbacks import Callback\n",
    "from sklearn.preprocessing import OneHotEncoder, StandardScaler\n",
    "from sklearn.metrics import confusion_matrix, f1_score, precision_score, recall_score\n",
    "\n",
    "# Essential Data Handling\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from math import ceil\n",
    "\n",
    "# Get Paths\n",
    "from glob import glob\n",
    "\n",
    "# EEG package\n",
    "from mne import pick_types\n",
    "from mne.io import read_raw_edf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CONTEXT = 'pn4/'\n",
    "MATERIAL = 'eegmmidb/'\n",
    "URL = 'https://www.physionet.org/' + CONTEXT + MATERIAL\n",
    "\n",
    "# Change this directory according to your setting\n",
    "USERDIR = './data/'\n",
    "\n",
    "page = requests.get(URL).text\n",
    "FOLDERS = sorted(list(set(re.findall(r'S[0-9]+', page))))\n",
    "\n",
    "URLS = [URL+x+'/' for x in FOLDERS]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Warning: Executing this block will create folders\n",
    "for folder in FOLDERS:\n",
    "    pathlib.Path(USERDIR +'/'+ folder).mkdir(parents=True, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Warning: Executing this block will start downloading data\n",
    "for i, folder in enumerate(FOLDERS):\n",
    "    page = requests.get(URLS[i]).text\n",
    "    subs = list(set(re.findall(r'S[0-9]+R[0-9]+', page)))\n",
    "    \n",
    "    print('Working on {}, {:.1%} completed'.format(folder, (i+1)/len(FOLDERS)))\n",
    "    for sub in subs:\n",
    "        urllib.request.urlretrieve(URLS[i]+sub+'.edf', os.path.join(USERDIR, folder, sub+'.edf'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Description\n",
    "\n",
    "Subjects performed different motor/imagery tasks while 64-channel EEG were recorded using the BCI2000 system (http://www.bci2000.org). Each subject performed 14 experimental runs: two one-minute baseline runs (one with eyes open, one with eyes closed), and three two-minute runs of each of the four following tasks:\n",
    "A target appears on either the left or the right side of the screen. The subject opens and closes the corresponding fist until the target disappears. Then the subject relaxes.\n",
    "A target appears on either the left or the right side of the screen. The subject imagines opening and closing the corresponding fist until the target disappears. Then the subject relaxes.\n",
    "A target appears on either the top or the bottom of the screen. The subject opens and closes either both fists (if the target is on top) or both feet (if the target is on the bottom) until the target disappears. Then the subject relaxes.\n",
    "A target appears on either the top or the bottom of the screen. The subject imagines opening and closing either both fists (if the target is on top) or both feet (if the target is on the bottom) until the target disappears. Then the subject relaxes.\n",
    "\n",
    "The data are provided here in EDF+ format (containing 64 EEG signals, each sampled at 160 samples per second, and an annotation channel). For use with PhysioToolkit software, rdedfann generated a separate PhysioBank-compatible annotation file (with the suffix .event) for each recording. The .event files and the annotation channels in the corresponding .edf files contain identical data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Raw Data Import\n",
    "\n",
    "I will use a EEG data handling package named MNE (https://martinos.org/mne/stable/index.html) to import raw data and annotation for events from edf files. This package also provides essential signal analysis features, e.g. band-pass filtering. The raw data were filtered using 1Hz of high-pass filter.\n",
    "\n",
    "In this research, there are 5 classes for the data: imagined motion of right fist, left fist, both fists, both feet, and rest with eyes closed. A data from one of the 109 subjects was excluded as the record was severely corrupted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get file paths\n",
    "PATH = './PhysioNet/'\n",
    "SUBS = glob(PATH + 'S[0-9]*')\n",
    "FNAMES = sorted([x[-4:] for x in SUBS])\n",
    "\n",
    "# Remove subject #89 with damaged data\n",
    "FNAMES.remove('S089')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_data(subj_num=FNAMES, epoch_sec=0.0625):\n",
    "    \"\"\" Import from edf files data and targets in the shape of 3D tensor\n",
    "    \n",
    "        Output shape: (Trial*Channel*TimeFrames)\n",
    "        \n",
    "        Some edf+ files recorded at low sampling rate, 128Hz, are excluded. \n",
    "        Majority was sampled at 160Hz.\n",
    "        \n",
    "        epoch_sec: time interval for one segment of mashes\n",
    "        \"\"\"\n",
    "    \n",
    "    # Event codes mean different actions for two groups of runs\n",
    "    run_type_0 = '02'.split(',')\n",
    "    run_type_1 = '04,08,12'.split(',')\n",
    "    run_type_2 = '06,10,14'.split(',')\n",
    "    \n",
    "    # Initiate X, y\n",
    "    X = []\n",
    "    y = []\n",
    "    \n",
    "    # To compute the completion rate\n",
    "    count = len(subj_num)\n",
    "    \n",
    "    # fixed numbers\n",
    "    nChan = 64 \n",
    "    sfreq = 160\n",
    "    sliding = epoch_sec/2 \n",
    "    timeFromQue = 0.5\n",
    "\n",
    "    # Sub-function to assign X and X, y\n",
    "    def append_X(n_segments, data, event=[]):\n",
    "        # Data should be changed\n",
    "        '''This function generate a tensor for X and append it to the existing X'''\n",
    "        \n",
    "        if len(event):\n",
    "            event_start = ceil(event[0] * sfreq)\n",
    "        else:\n",
    "            event_start = 0\n",
    "    \n",
    "        def window(n):\n",
    "            windowStart = int(timeFromQue*sfreq) + int(sfreq*sliding*n) + event_start\n",
    "            windowEnd = int(timeFromQue*sfreq) + int(sfreq*sliding*(n+2)) + event_start\n",
    "            \n",
    "            while (windowEnd - windowStart) != 10:\n",
    "                windowEnd += int(sfreq*epoch_sec) - (windowEnd - windowStart)\n",
    "                \n",
    "            return [windowStart, windowEnd]\n",
    "        \n",
    "        new_x = [data[:, window(n)[0]: window(n)[1]] for n in range(n_segments)\\\n",
    "                 if data[:, window(n)[0]:window(n)[1]].shape==(nChan, int(sfreq*epoch_sec))]\n",
    "        return new_x\n",
    "    \n",
    "    def append_X_Y(run_type, event, old_x, old_y, data):\n",
    "        '''This function seperate the type of events \n",
    "        (refer to the data descriptitons for the list of the types)\n",
    "        Then assign X and Y according to the event types'''\n",
    "        # Number of sliding windows\n",
    "        n_segments = int(event[1]/epoch_sec)\n",
    "        \n",
    "        # Rest excluded\n",
    "        if event[2] == 'T0':\n",
    "            return old_x, old_y\n",
    "        \n",
    "        # y assignment\n",
    "        if run_type == 1:\n",
    "            temp_y = [1] if event[2] == 'T1' else [2]\n",
    "        \n",
    "        elif run_type == 2:\n",
    "            temp_y = [3] if event[2] == 'T1' else [4]\n",
    "                \n",
    "        new_x = append_X(n_segments, data, event)\n",
    "        new_y = old_y + temp_y*len(new_x)\n",
    "        \n",
    "        return old_x + new_x, new_y\n",
    "    \n",
    "    # Iterate over subj_num: S001, S002, S003...\n",
    "    for i, subj in enumerate(subj_num):\n",
    "        # Return completion rate\n",
    "        if i%((len(subj_num)//10)+1) == 0:\n",
    "            print('working on {}, {:.0%} completed'.format(subj, i/count))\n",
    "\n",
    "        # Get file names\n",
    "        fnames = glob(os.path.join(PATH, subj, subj+'R*.edf'))\n",
    "        fnames = sorted([name for name in fnames if name[-6:-4] in run_type_0+run_type_1+run_type_2])\n",
    "        \n",
    "        for i, fname in enumerate(fnames):\n",
    "            \n",
    "            # Import data into MNE raw object\n",
    "            raw = read_raw_edf(fname, preload=True, verbose=False)\n",
    "\n",
    "            picks = pick_types(raw.info, eeg=True)\n",
    "            \n",
    "            if raw.info['sfreq'] != 160:\n",
    "                print('{} is sampled at 128Hz so will be excluded.'.format(subj))\n",
    "                break\n",
    "            \n",
    "            # High-pass filtering\n",
    "            raw.filter(l_freq=1, h_freq=None, picks=picks)\n",
    "            \n",
    "            # Get annotation\n",
    "            try:\n",
    "                events = raw.find_edf_events()\n",
    "            except:\n",
    "                continue\n",
    "                \n",
    "            # Get data\n",
    "            data = raw.get_data(picks=picks)\n",
    "            \n",
    "            # Number of this run\n",
    "            which_run = fname[-6:-4]\n",
    "            \n",
    "            \"\"\" Assignment Starts \"\"\" \n",
    "            # run 1 - baseline (eye closed)\n",
    "            if which_run in run_type_0:\n",
    "                \n",
    "                # Number of sliding windows\n",
    "                n_segments = int((raw.n_times/(epoch_sec*sfreq)))\n",
    "\n",
    "                # Append 0`s based on number of windows\n",
    "                new_X = append_X(n_segments, data)\n",
    "                X += new_X\n",
    "                y.extend([0] * len(new_X))\n",
    "                    \n",
    "            # run 4,8,12 - imagine opening and closing left or right fist    \n",
    "            elif which_run in run_type_1:\n",
    "                \n",
    "                for i, event in enumerate(events):\n",
    "                    X, y = append_X_Y(run_type=1, event=event, old_x=X, old_y=y, data=data)\n",
    "                        \n",
    "            # run 6,10,14 - imagine opening and closing both fists or both feet\n",
    "            elif which_run in run_type_2:\n",
    "                   \n",
    "                for i, event in enumerate(events):         \n",
    "                    X, y = append_X_Y(run_type=2, event=event, old_x=X, old_y=y, data=data)\n",
    "                        \n",
    "    X = np.stack(X)\n",
    "    y = np.array(y).reshape((-1,1))\n",
    "    return X, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "X,y = get_data(FNAMES, epoch_sec=0.0625)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(X.shape)\n",
    "print(y.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Data Preprocessing\n",
    "\n",
    "The original goal of applying neural networks is to exclude hand-crafted algorithms & preprocessing as much as possible. I did not use any proprecessing techniques further than standardization to build an end-to-end classifer from the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.preprocessing import OneHotEncoder, scale\n",
    "\n",
    "#%%\n",
    "def convert_mesh(X):\n",
    "    \n",
    "    mesh = np.zeros((X.shape[0], X.shape[2], 10, 11, 1))\n",
    "    X = np.swapaxes(X, 1, 2)\n",
    "    \n",
    "    # 1st line\n",
    "    mesh[:, :, 0, 4:7, 0] = X[:,:,21:24]; print('1st finished')\n",
    "    \n",
    "    # 2nd line\n",
    "    mesh[:, :, 1, 3:8, 0] = X[:,:,24:29]; print('2nd finished')\n",
    "    \n",
    "    # 3rd line\n",
    "    mesh[:, :, 2, 1:10, 0] = X[:,:,29:38]; print('3rd finished')\n",
    "    \n",
    "    # 4th line\n",
    "    mesh[:, :, 3, 1:10, 0] = np.concatenate((X[:,:,38].reshape(-1, X.shape[1], 1),\\\n",
    "                                          X[:,:,0:7], X[:,:,39].reshape(-1, X.shape[1], 1)), axis=2)\n",
    "    print('4th finished')\n",
    "    \n",
    "    # 5th line\n",
    "    mesh[:, :, 4, 0:11, 0] = np.concatenate((X[:,:,(42, 40)],\\\n",
    "                                        X[:,:,7:14], X[:,:,(41, 43)]), axis=2)\n",
    "    print('5th finished')\n",
    "    \n",
    "    # 6th line\n",
    "    mesh[:, :, 5, 1:10, 0] = np.concatenate((X[:,:,44].reshape(-1, X.shape[1], 1),\\\n",
    "                                        X[:,:,14:21], X[:,:,45].reshape(-1, X.shape[1], 1)), axis=2)\n",
    "    print('6th finished')\n",
    "               \n",
    "    # 7th line\n",
    "    mesh[:, :, 6, 1:10, 0] = X[:,:,46:55]; print('7th finished')\n",
    "    \n",
    "    # 8th line\n",
    "    mesh[:, :, 7, 3:8, 0] = X[:,:,55:60]; print('8th finished')\n",
    "    \n",
    "    # 9th line\n",
    "    mesh[:, :, 8, 4:7, 0] = X[:,:,60:63]; print('9th finished')\n",
    "    \n",
    "    # 10th line\n",
    "    mesh[:, :, 9, 5, 0] = X[:,:,63]; print('10th finished')\n",
    "    \n",
    "    return mesh\n",
    "\n",
    "#%%\n",
    "def prepare_data(X, y, test_ratio=0.2, return_mesh=True, set_seed=42):\n",
    "    \n",
    "    # y encoding\n",
    "    oh = OneHotEncoder()\n",
    "    y = oh.fit_transform(y).toarray()\n",
    "    \n",
    "    # Shuffle trials\n",
    "    np.random.seed(set_seed)\n",
    "    trials = X.shape[0]\n",
    "    shuffle_indices = np.random.permutation(trials)\n",
    "    X = X[shuffle_indices]\n",
    "    y = y[shuffle_indices]\n",
    "    \n",
    "    # Test set seperation\n",
    "    train_size = int(trials*(1-test_ratio)) \n",
    "    X_train, X_test, y_train, y_test = X[:train_size,:,:], X[train_size:,:,:],\\\n",
    "                                    y[:train_size,:], y[train_size:,:]\n",
    "                                    \n",
    "    # Z-score Normalization\n",
    "    def scale_data(X):\n",
    "        shape = X.shape\n",
    "        for i in range(shape[0]):\n",
    "            X[i,:, :] = scale(X[i,:, :])\n",
    "            if i%int(shape[0]//10) == 0:\n",
    "                print('{:.0%} done'.format((i+1)/shape[0]))   \n",
    "        return X\n",
    "            \n",
    "    X_train, X_test  = scale_data(X_train), scale_data(X_test)\n",
    "    if return_mesh:\n",
    "        X_train, X_test = convert_mesh(X_train), convert_mesh(X_test)\n",
    "    \n",
    "    return X_train, y_train, X_test, y_test\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, y_train, X_test, y_test = prepare_data(X, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As the EEG recording instrument has 3D locations over the subjects\\` scalp, it is essential for the model to learn from the spatial pattern as well as the temporal pattern. I transformed the data into 2D meshes that represents the locations of the electrodes so that stacked convolutional neural networks can grasp the spatial information."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Modeling - Time-Distributed CNN + RNN\n",
    "\n",
    "Training Plan:\n",
    "\n",
    "+ 4 GPU units (Nvidia Tesla P100) were used to train this neural network.\n",
    "+ Instead of training the whole model at once, I trained the first block (CNN) first. Then using the trained parameters as initial values, I trained the next blocks step-by-step. This approach can greatly reduce the time required for training and help avoiding falling into local minimums.\n",
    "+ The first blocks (CNN) can be applied for other EEG classification models as a pre-trained base.\n",
    "\n",
    "+ The initial learning rate is set to be $10^{3}$ with Adam optimization. I used several callbacks such as ReduceLROnPlateau which adjusts the learning rate at local minima. Also, I record the log for tensorboard to monitor the training process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = X_train.squeeze().reshape(*X_train.squeeze().shape, 1)\n",
    "X_test = X_test.squeeze().reshape(*X_test.squeeze().shape, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make another dimension, 1, to apply CNN for each time frame.\n",
    "X_train = X_train.reshape(*X_train.shape, 1)\n",
    "X_test = X_test.reshape(*X_test.shape, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 Keras Implementation\n",
    "\n",
    "The Keras functional API is the way to go for defining complex models, such as multi-output models, directed acyclic graphs, or models with shared layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "## Complicated Model - the same as Zhang`s\n",
    "input_shape = (10, 10, 11, 1)\n",
    "lecun = initializers.lecun_normal(seed=42)\n",
    "\n",
    "# TimeDistributed Wrapper\n",
    "def timeDist(layer, prev_layer, name):\n",
    "    return TimeDistributed(layer, name=name)(prev_layer)\n",
    "    \n",
    "# Input layer\n",
    "inputs = Input(shape=input_shape)\n",
    "\n",
    "# Convolutional layers block\n",
    "x = timeDist(Conv2D(32, (3,3), padding='same', \n",
    "                    data_format='channels_last', kernel_initializer=lecun), inputs, name='CNN1')\n",
    "x = BatchNormalization(name='batch1')(x)\n",
    "x = Activation('elu', name='act1')(x)\n",
    "x = timeDist(Conv2D(64, (3,3), padding='same', data_format='channels_last', kernel_initializer=lecun), x, name='CNN2')\n",
    "x = BatchNormalization(name='batch2')(x)\n",
    "x = Activation('elu', name='act2')(x)\n",
    "x = timeDist(Conv2D(128, (3,3), padding='same', data_format='channels_last', kernel_initializer=lecun), x, name='CNN3')\n",
    "x = BatchNormalization(name='batch3')(x)\n",
    "x = Activation('elu', name='act3')(x)\n",
    "x = timeDist(Flatten(), x, name='flatten')\n",
    "\n",
    "# Fully connected layer block\n",
    "y = Dense(1024, kernel_initializer=lecun, name='FC')(x)\n",
    "y = Dropout(0.5, name='dropout1')(y)\n",
    "y = BatchNormalization(name='batch4')(y)\n",
    "y = Activation(activation='elu')(y)\n",
    "\n",
    "# Recurrent layers block\n",
    "z = LSTM(64, kernel_initializer=lecun, return_sequences=True, name='LSTM1')(y)\n",
    "z = LSTM(64, kernel_initializer=lecun, name='LSTM2')(z)\n",
    "\n",
    "# Fully connected layer block\n",
    "h = Dense(1024, kernel_initializer=lecun, activation='elu', name='FC2')(z)\n",
    "h = Dropout(0.5, name='dropout2')(h)\n",
    "\n",
    "# Output layer\n",
    "outputs = Dense(5, activation='softmax')(h)\n",
    "\n",
    "# Model compile\n",
    "model = Model(inputs=inputs, outputs=outputs)\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "model = load_model('./model_1230.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "\n",
    "# Load a model to transfer pre-trained parameters\n",
    "trans_model = model.load('CNN_3blocks.h5')\n",
    "\n",
    "# Transfer learning - parameter copy & paste\n",
    "which_layer = 'CNN1,CNN2,CNN3,batch1,batch2,batch3'.split(',')\n",
    "layer_names = [layer.name for layer in model.layers]\n",
    "trans_layer_names = [layer.name for layer in trans_model.layers]\n",
    "\n",
    "for layer in which_layer:\n",
    "    ind = layer_names.index(layer)\n",
    "    trans_ind = trans_layer_names.index(layer)\n",
    "    model.layers[ind].set_weights(trans_model.layers[trans_ind].get_weights())\n",
    "    \n",
    "for layer in model.layers[:9]: # Freeze the first 9 layers(CNN block)\n",
    "    layer.trainable = False\n",
    "    \n",
    "    \n",
    "# Turn on multi-GPU mode\n",
    "model = multi_gpu_model(model, gpus=4)\n",
    "\n",
    "\n",
    "This metrics calculate sensitivity and specificity batch-wise.\n",
    "Keras development team removed this feature because\n",
    "these metrics should be understood as global metrics.\n",
    "\n",
    "I am not using it this time.\n",
    "\n",
    "# Metrics - sensitivity, specificity, accuracy\n",
    "def sens(y_true, y_pred): # Sensitivity\n",
    "    true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))\n",
    "    possible_positives = K.sum(K.round(K.clip(y_true, 0, 1)))\n",
    "    return true_positives / (possible_positives + K.epsilon())\n",
    "\n",
    "def prec(y_true, y_pred): # Precision\n",
    "    true_negatives = K.sum(K.round(K.clip((1-y_true) * (1-y_pred), 0, 1)))\n",
    "    possible_negatives = K.sum(K.round(K.clip(1-y_true, 0, 1)))\n",
    "    return true_negatives / (possible_negatives + K.epsilon())\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from keras.callbacks import Callback\n",
    "\n",
    "class CustomModelCheckPoint(Callback):\n",
    "    def __init__(self,**kargs):\n",
    "        super(CustomModelCheckPoint,self).__init__(**kargs)\n",
    "        self.epoch_accuracy = {} # loss at given epoch\n",
    "        self.epoch_loss = {} # accuracy at given epoch\n",
    "        def on_epoch_begin(self,epoch, logs={}):\n",
    "            # Things done on beginning of epoch. \n",
    "            return\n",
    "\n",
    "        def on_epoch_end(self, epoch, logs={}):\n",
    "            # things done on end of the epoch\n",
    "            self.epoch_accuracy[epoch] = logs.get(\"acc\")\n",
    "            self.epoch_loss[epoch] = logs.get(\"loss\")\n",
    "            #self.model.save_weights(\"name-of-model-%d.h5\" %epoch) # save the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "callbacks_list = [callbacks.ModelCheckpoint('model_1230.h5', save_best_only=True, monitor='val_loss'),\n",
    "                 callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=5),\n",
    "                 callbacks.CSVLogger('model_1230_res.csv', separator=',', append=True)]\n",
    "\n",
    "# Start training\n",
    "model.compile(loss='categorical_crossentropy', optimizer=optimizers.adam(lr=1e-4), metrics=['acc'])\n",
    "history = model.fit(X_train, y_train, batch_size=64, epochs=5000, shuffle=True,\n",
    "                    validation_split=0.2, callbacks=callbacks_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2 Tensorflow Eager Execution API\n",
    "\n",
    "TensorFlow's eager execution is an imperative programming environment that evaluates operations immediately, without building graphs: operations return concrete values instead of constructing a computational graph to run later. This makes it easy to get started with TensorFlow and debug models, and it reduces boilerplate as well. To follow along with this guide, run the code samples below in an interactive python interpreter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parameters\n",
    "learning_rate = 0.001\n",
    "num_steps = 1000\n",
    "batch_size = 128\n",
    "display_step = 100\n",
    "\n",
    "# Network Parameters\n",
    "num_input = X.shape[0] # PhysioNet data input (mesh shape: 10*11)\n",
    "num_classes = 5 # PhysioNet total classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using TF Dataset to split data into batches\n",
    "dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train)).batch(batch_size)\n",
    "dataset_iter = tfe.Iterator(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ZhangModel(tf.keras.Model):\n",
    "    def __init__(self, n_nodes=[3,2,2], \n",
    "                 initializer=tf.contrib.layers.variance_scaling_initializer(mode=\"FAN_AVG\")):\n",
    "        \"\"\"\n",
    "        This is a tensorflow implementation of Zhang`s model (2018) with \n",
    "        3 CNN, 2 LSTM, 2 dense layers by default.\n",
    "        \n",
    "        n_nodes [list] : specifies the number of layers for each block. \n",
    "                        [CNN, LSTM, DENSE] respectively.\n",
    "        initializer [tf.layers] : defualt initializer set to be He initializer\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.n_CNN, self.n_LSTM, self.n_dense = n_nodes\n",
    "        self.CNN, self.LSTM, self.dense = [[] for i in range(len(n_nodes))]\n",
    "        self.initializer = initializer\n",
    "        \n",
    "        count = 0\n",
    "        for n in range(self.n_CNN):\n",
    "            count += 1\n",
    "            n_filter = 32*2**count\n",
    "            \n",
    "            self.CNN.append(tf.keras.layers.Conv2D(n_filter, (3, 3), padding='same', \n",
    "                                                data_format='channels_last',\n",
    "                                                kernel_initializer=self.initializer))\n",
    "        \n",
    "        for n in range(self.n_LSTM):    \n",
    "            n_hidden = 64\n",
    "            return_sequences = False if n == self.n_LSTM-1 else True\n",
    "            \n",
    "            name = 'LSTM' + str(len(self.LSTM)+1)            \n",
    "            self.LSTM.append(tf.keras.layers.LSTM(n_hidden, kernel_initializer=self.initializer, \n",
    "                                             return_sequences=return_sequences, name=name))\n",
    "            \n",
    "        for n in range(self.n_dense):\n",
    "            n_node=1024\n",
    "            \n",
    "            name = 'Dense' + str(len(self.dense)+1)\n",
    "            self.dense.append(tf.keras.layers.Dense(n_node, \n",
    "                                                    kernel_initializer=self.initializer, name=name))\n",
    "\n",
    "    def call(self, input_tensor):\n",
    "        \"Run the model.\"\n",
    "        \n",
    "        assert self.CNN, 'No CNN blocks defined!'\n",
    "        assert self.dense, 'No Dense Blocks defined!'\n",
    "        assert self.LSTM, 'No LSTM blocks defined!'\n",
    "        \n",
    "        def timeDist(layer, prev_layer, name):\n",
    "            return tf.keras.layers.TimeDistributed(layer, name=name)(prev_layer)\n",
    "        \n",
    "        for i, layer in enumerate(self.CNN):\n",
    "            name = 'CNN' + str(i+1)\n",
    "            nameBatch = 'batch' + str(i+1)\n",
    "            nameAct = 'act' + str(i+1)\n",
    "            \n",
    "            prev_layer = input_tensor if i==0 else x\n",
    "            \n",
    "            x = timeDist(layer, prev_layer, name=name)\n",
    "            x = tf.keras.layers.BatchNormalization(name=nameBatch)(x)\n",
    "            x = tf.nn.elu(x, name=nameAct)\n",
    "            \n",
    "        x = timeDist(tf.keras.layers.Flatten(), x, name='flatten')\n",
    "        \n",
    "        for i, layer in enumerate(self.dense):            \n",
    "            \n",
    "            nameDrop = 'drop' + str(i+1)\n",
    "            nameBatch = 'batch' + str(i+1)\n",
    "            nameAct = 'act' + str(i+1)\n",
    "            \n",
    "            if i == len(self.dense)-1:\n",
    "                break\n",
    "            \n",
    "            x = layer(x)\n",
    "            x = tf.keras.layers.Dropout(0.5, name=nameDrop)(x)\n",
    "            x = tf.keras.layers.BatchNormalization(name=nameBatch)(x)\n",
    "            x = tf.nn.elu(x, name=nameAct)\n",
    "            \n",
    "        for i, layer in enumerate(self.LSTM):\n",
    "            x = layer(x)\n",
    "        \n",
    "        x = self.dense[-1](x)\n",
    "        x = tf.keras.layers.Dropout(0.5, name=nameDrop)(x)\n",
    "        x = tf.keras.layers.BatchNormalization(name=nameBatch)(x)\n",
    "        x = tf.nn.elu(x, name=nameAct)\n",
    "        \n",
    "        output = tf.keras.layers.Dense(5, activation='softmax')(x)\n",
    "        \n",
    "        return output    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model = ZhangModel([3, 2, 2])\n",
    "print(model(tf.random_normal([1, 10, 10, 11, 1])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cross-Entropy loss function\n",
    "def loss_fn(inference_fn, inputs, labels):\n",
    "    # Using sparse_softmax cross entropy\n",
    "    return tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(\n",
    "        logits=inference_fn(inputs), labels=labels))\n",
    "\n",
    "# Calculate accuracy\n",
    "def accuracy_fn(inference_fn, inputs, labels):\n",
    "    prediction = inference_fn(inputs)\n",
    "    correct_pred = tf.equal(tf.argmax(prediction, 1), tf.argmax(labels, 1))\n",
    "    return tf.reduce_mean(tf.cast(correct_pred, tf.float32))\n",
    "\n",
    "# SGD Optimizer\n",
    "optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)\n",
    "\n",
    "# Compute gradients\n",
    "grad = tfe.implicit_gradients(loss_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training\n",
    "average_loss = 0.\n",
    "average_acc = 0.\n",
    "for step in range(num_steps):\n",
    "\n",
    "    # Iterate through the dataset\n",
    "    try:\n",
    "        d = dataset_iter.next()\n",
    "    except StopIteration:\n",
    "        # Refill queue\n",
    "        dataset_iter = tfe.Iterator(dataset)\n",
    "        d = dataset_iter.next()\n",
    "\n",
    "    # EEGs\n",
    "    x_batch = d[0]\n",
    "    # Labels\n",
    "    y_batch = tf.cast(d[1], dtype=tf.int64)\n",
    "\n",
    "    # Compute the batch loss\n",
    "    batch_loss = loss_fn(model, x_batch, y_batch)\n",
    "    average_loss += batch_loss\n",
    "    \n",
    "    # Compute the batch accuracy\n",
    "    batch_accuracy = accuracy_fn(model, x_batch, y_batch)\n",
    "    average_acc += batch_accuracy\n",
    "\n",
    "    if step == 0:\n",
    "        # Display the initial cost, before optimizing\n",
    "        print(\"Initial loss= {:.9f}\".format(average_loss))\n",
    "\n",
    "    # Update the variables following gradients info\n",
    "    optimizer.apply_gradients(grad(model, x_batch, y_batch))\n",
    "\n",
    "    # Display info\n",
    "    if (step + 1) % display_step == 0 or step == 0:\n",
    "        if step > 0:\n",
    "            average_loss /= display_step\n",
    "            average_acc /= display_step\n",
    "        print(\"Step:\", '%04d' % (step + 1), \" loss=\",\n",
    "              \"{:.9f}\".format(average_loss), \" accuracy=\",\n",
    "              \"{:.4f}\".format(average_acc))\n",
    "        average_loss = 0.\n",
    "        average_acc = 0."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load in libraries\n",
    "import pickle\n",
    "import itertools\n",
    "from sklearn.metrics import confusion_matrix, classification_report, accuracy_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make directories\n",
    "if not os.path.exists('./metrics/'):\n",
    "    os.makedirs('./metrics/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_history(history):\n",
    "    loss_list = [s for s in history.keys() if 'loss' in s and 'val' not in s]\n",
    "    val_loss_list = [s for s in history.keys() if 'loss' in s and 'val' in s]\n",
    "    acc_list = [s for s in history.keys() if 'acc' in s and 'val' not in s]\n",
    "    val_acc_list = [s for s in history.keys() if 'acc' in s and 'val' in s]\n",
    "    \n",
    "    if len(loss_list) == 0:\n",
    "        print('Loss is missing in history')\n",
    "        return \n",
    "    \n",
    "    ## As loss always exists\n",
    "    epochs = range(1,len(history[loss_list[0]]) + 1)\n",
    "    \n",
    "   ## Loss\n",
    "    plt.figure(1)\n",
    "    for l in loss_list:\n",
    "        plt.plot(epochs, history[l], 'b', label='Training loss (' + str(str(format(history[l][-1],'.5f'))+')'))\n",
    "    for l in val_loss_list:\n",
    "        plt.plot(epochs, history[l], 'g', label='Validation loss (' + str(str(format(history[l][-1],'.5f'))+')'))\n",
    "    \n",
    "    plt.title('Loss')\n",
    "    plt.xlabel('Epochs')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.legend()\n",
    "    plt.savefig(\"./metrics/loss.png\")\n",
    "    \n",
    "    ## Accuracy\n",
    "    plt.figure(2)\n",
    "    for l in acc_list:\n",
    "        plt.plot(epochs, history[l], 'b', label='Training accuracy (' + str(format(history[l][-1],'.5f'))+')')\n",
    "    for l in val_acc_list:    \n",
    "        plt.plot(epochs, history[l], 'g', label='Validation accuracy (' + str(format(history[l][-1],'.5f'))+')')\n",
    "\n",
    "    plt.title('Accuracy')\n",
    "    plt.xlabel('Epochs')\n",
    "    plt.ylabel('Accuracy')\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "    plt.savefig(\"./metrics/acc.png\")\n",
    "    \n",
    "def plot_confusion_matrix(cm, classes,\n",
    "                          normalize=False,\n",
    "                          cmap=plt.cm.Blues):\n",
    "    \"\"\"\n",
    "    This function prints and plots the confusion matrix.\n",
    "    Normalization can be applied by setting `normalize=True`.\n",
    "    \"\"\"\n",
    "    if normalize:\n",
    "        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]\n",
    "        title='Normalized confusion matrix'\n",
    "    else:\n",
    "        title='Confusion matrix'\n",
    "\n",
    "    plt.figure(3)\n",
    "    plt.imshow(cm, interpolation='nearest', cmap=cmap)\n",
    "    plt.title(title)\n",
    "    plt.colorbar()\n",
    "    tick_marks = np.arange(len(classes))\n",
    "    plt.xticks(tick_marks, classes, rotation=45)\n",
    "    plt.yticks(tick_marks, classes)\n",
    "\n",
    "    fmt = '.2f' if normalize else 'd'\n",
    "    thresh = cm.max() / 2.\n",
    "    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):\n",
    "        plt.text(j, i, format(cm[i, j], fmt),\n",
    "                 horizontalalignment=\"center\",\n",
    "                 color=\"white\" if cm[i, j] > thresh else \"black\")\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.ylabel('True label')\n",
    "    plt.xlabel('Predicted label')\n",
    "    plt.savefig(\"./metrics/confuMat.png\")\n",
    "    plt.show()\n",
    "    \n",
    "def full_multiclass_report(model,\n",
    "                           x,\n",
    "                           y_true,\n",
    "                           classes):\n",
    "    \n",
    "    # 2. Predict classes and stores in y_pred\n",
    "    y_pred = model.predict(x).argmax(axis=1)\n",
    "    \n",
    "    # 3. Print accuracy score\n",
    "    print(\"Accuracy : \"+ str(accuracy_score(y_true,y_pred)))\n",
    "    \n",
    "    print(\"\")\n",
    "    \n",
    "    # 4. Print classification report\n",
    "    print(\"Classification Report\")\n",
    "    print(classification_report(y_true,y_pred,digits=4))    \n",
    "    \n",
    "    # 5. Plot confusion matrix\n",
    "    cnf_matrix = confusion_matrix(y_true,y_pred)\n",
    "    print(cnf_matrix)\n",
    "    plot_confusion_matrix(cnf_matrix,classes=classes)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load in the data\n",
    "howManyTest = 0.2\n",
    "\n",
    "thisInd = np.random.randint(0, len(X_test), size=(len(X_test)//howManyTest))\n",
    "X_conf, y_conf = X_test[[i for i in thisInd], :], y_test[[i for i in thisInd],:] \n",
    "\n",
    "'''\n",
    "## Only if you have a previous model + history\n",
    "# Get the model\n",
    "model = models.load_model('./model/model0.h5')\n",
    "\n",
    "# Get the history\n",
    "with open('./history/history0.pkl', 'rb') as hist:\n",
    "    history = pickle.load(hist)\n",
    "'''\n",
    "\n",
    "# Get the graphics\n",
    "plot_history(history)\n",
    "X_test = X_test.reshape(X_test.shape[0], X_train.shape[1], X_train.shape[2], X_train.shape[3], 1)\n",
    "full_multiclass_report(model,\n",
    "                       X_test,\n",
    "                       y_test.argmax(axis=1),\n",
    "                       [1,2,3,4,5])"
   ]
  }
 ],
 "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}