EEG-LSTM-01.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Sep 30 14:42:30 2019

@author: jeremiah
"""
#%% inportings
import codecs
import arff as af
import torch
import numpy as np
from torch import nn, optim
import matplotlib.pyplot as plt
%matplotlib inline


#%% loading data

path = '/home/jeremiah/data_analysis/datasets/EEG Eye State.arff'

file = codecs.open(path, 'rb', 'utf-8')
file = af.load(file)

#print(af.dumps(file))

#%% returns training, validation and testing datasets along with training, validation and testing labels

def get_labels(data, test_len_per, val_len_per):

    total_len = len(data)
    train_len = int((1-val_len_per)*total_len)
    train_labels, val_labels = data[:train_len], data[train_len:]

    #divide train_data between train_data and test_data
    total_len = len(train_labels)
    train_len = int(0.7*total_len)
    train_labels, test_labels = train_labels[:train_len], train_labels[train_len:]

    return train_labels, test_labels, val_labels


def divide_dataset(data, test_len_per, val_len_per):
    #divide data between train_data and val_data

    labels = []

    for lista in data:
        labels.append(lista[-1])
        lista.pop()


    total_len = len(data)
    train_len = int((1-val_len_per)*total_len)
    train_data, val_data = data[:train_len], data[train_len:]

    #divide train_data between train_data and test_data
    total_len = len(train_data)
    train_len = int(0.7*total_len)
    train_data, test_data = train_data[:train_len], train_data[train_len:]


    train_labels, test_labels, val_labels = get_labels(labels, test_len_per, val_len_per)

    return train_data, train_labels, test_data, test_labels, val_data, val_labels

#%% dividing data between training, validation and testing sets


train_data, train_labels, test_data, test_labels, val_data, val_labels = divide_dataset(file['data'], 0.3, 0.1)

train_data =  np.asarray(train_data, dtype='float32')
train_labels =  np.asarray(train_labels, dtype='float32')

test_data =  np.asarray(test_data, dtype='float32')
test_labels =  np.asarray(test_labels, dtype='float32')

val_data =  np.asarray(val_data, dtype='float32')
val_labels =  np.asarray(val_labels, dtype='float32')

#%%
print(f'training set length: {len(train_data)}')
print(f'testing set length: {len(test_data)}')
print(f'validation set length: {len(val_data)}')

#%%

def get_batches(inputs, labels, batch_size, seq_len):

    batch_total_size = batch_size * seq_len
    n_batches = len(inputs)//batch_total_size

    inputs = inputs[:n_batches * batch_total_size]
    inputs = inputs.reshape((batch_size, -1, 14))

    labels = labels[:n_batches * batch_total_size]
    labels = labels.reshape((batch_size, -1))

    for i in range(0, inputs.shape[1], seq_len):

        x = inputs[:, i : i+seq_len]
        y = labels[:, i : i+seq_len]

        yield x, y

#%% defining our model

class LSTM_Net(nn.Module):
    def __init__(self, in_size, out_size, h_size, n_layers, drop_p = 0.5, batch_first = True):
        super().__init__()

        self.in_size = in_size
        self.out_size = out_size
        self.hidden_size = h_size
        self.n_layers = n_layers

        self.lstm = nn.LSTM(in_size, h_size, n_layers, dropout = drop_p, batch_first = batch_first)

        self.dropout = nn.Dropout(drop_p)

        self.fc = nn.Linear(h_size, out_size)
    def forward(self, x, h):
        out, h = self.lstm(x, h)

        out = self.dropout(out)

        out = out.reshape(-1, self.hidden_size)

        out = self.fc(out)
        return out, h

    def init_hidden(self, batch_size):
        weight = next(self.parameters()).data

        hidden = (weight.new(self.n_layers, batch_size, self.hidden_size).zero_(),
                  weight.new(self.n_layers, batch_size, self.hidden_size).zero_())

        return hidden


#%% training the model

def train(net, inputs, labels, test_inputs, test_labels, epochs, batch_size, seq_len, print_every, clip = 5, lr=0.001):

    net.train()

    criterion = nn.MSELoss()
    op = optim.Adam(net.parameters(), lr=lr)

    steps = 0
    last_test_loss = np.Inf
    train_losses = []
    test_losses = []
    model_name = ''

    for e in range(epochs):

        hidden = net.init_hidden(batch_size)
        test_loss = []
        train_loss = []
        for x, y in get_batches(inputs, labels, batch_size, seq_len):

            steps += 1

            x, y = torch.from_numpy(x), torch.from_numpy(y)

            hidden = tuple([each.data for each in hidden])

            net.zero_grad()
            out, hidden = net(x, hidden)
            #print(out.shape)

            loss = criterion(out, y.reshape(batch_size*seq_len, 1))
            train_loss.append(loss.item())

            loss.backward()
            nn.utils.clip_grad_norm_(net.parameters(), clip)
            op.step()

        else:

            net.eval()

            test_loss = []
            test_h = net.init_hidden(batch_size)

            for x, y in get_batches(test_inputs, test_labels, batch_size, seq_len):

                x, y = torch.from_numpy(x), torch.from_numpy(y)

                test_h = tuple([each.data for each in test_h])

                out, test_h = net(x, test_h)

                test_loss.append(criterion(out, y.reshape(batch_size*seq_len, 1)).item())


            net.train()
            train_losses.append(np.mean(train_loss))
            test_losses.append(np.mean(test_loss))

            print('Epoch: {}/{}'.format(e+1, epochs),
                  'Step: {}'.format(steps),
                  'Train_Loss: {}'.format(train_losses[-1]),
                  'Test_Loss: {}'.format(test_losses[-1]))

            if test_losses[-1] < last_test_loss:

                print('Saving model...')
                last_test_loss = test_losses[-1]
                model_name = '/home/jeremiah/trained_models/LSTM-EGG-01-' + str(epochs) + '-epochs-' + str(net.hidden_size) + '-h_layers-bs' + str(batch_size) +'-sl' + str(seq_len) + '.net'

                checkpoint = {'input_size' : net.in_size,
                              'output_size' : net.out_size,
                              'hidden_size' : net.hidden_size,
                              'n_layers' : net.n_layers,
                              'state_dict' : net.state_dict()}
                with open(model_name, 'wb') as f:
                    torch.save(checkpoint, f)

                model_name = '/home/jeremiah/trained_models/LSTM-EGG-01-' + str(epochs) + '-epochs-' + str(net.hidden_size) + '-h_layers-bs' + str(batch_size) +'-sl' + str(seq_len) + '.txt'
                net_info = 'Model name: ' + model_name[:-4] + '\n'
                net_info += 'Steps: ' + str(steps) + '\n'
                net_info += 'Epochs: ' + str(epochs) + '\n'
                net_info += 'Batch size: ' + str(batch_size) + '\n'
                net_info += 'Sequence length: ' + str(seq_len) + '\n'
                net_info += 'Learning rate: ' + str(learn_rate) + '\n\n'
                net_info += 'Hidden size: ' + str(net.hidden_size) + '\n'
                net_info += 'Number of layers: ' + str(net.n_layers) + '\n\n'
                net_info += 'Training loss: ' + str(train_losses[-1]) + '\n'
                net_info += 'Testing loss: ' + str(test_losses[-1])+'\n'
                net_info += 'Model saved in epoch: ' + str(e) + '\n\n'
                net_info += 'Training set length: ' + str(len(train_data)) + '\n'
                net_info += 'Testing set length: ' + str(len(test_data)) + '\n'
                net_info += 'Validation set lengt: ' + str(len(val_data)) + '\n\n'
                net_info += 'Loss function: MSELoss() - Mean Squared Error Loss Function\n'

                with open(model_name, 'w+') as f:
                    f.write(net_info)

    plt.plot(train_losses, 'r', label='train_loss')
    plt.plot(test_losses, 'b', label='test_loss')
    plt.legend(loc='best')
    plt.savefig('/home/jeremiah/trained_models/LSTM-EEG-01-'+str(epochs)+'-epochs-'+str(hidden_size)+'-h_layers-bs' + str(batch_size) +'-sl' + str(seq_len) + '.png')



#%%
input_size = 14
output_size = 1
hidden_size = 32
n_layers = 2
dropout = 0.5


net = LSTM_Net(input_size, output_size, hidden_size, n_layers, dropout, batch_first=True)
print(net)

#%%

epochs = 100
batch_size = 16
seq_length = 8
print_every = 100
clip = 5
learn_rate = 0.002

train(net, train_data, train_labels, test_data, test_labels, epochs, batch_size, seq_length, print_every, clip, learn_rate)

#%% load a trained model
with open('/home/jeremiah/trained_models/LSTM-EGG-01-70-epochs-16-h_layers-bs4-sl1.net', 'rb') as f:
    cpt = torch.load(f)

net = LSTM_Net(cpt['input_size'], cpt['output_size'], cpt['hidden_size'], cpt['n_layers'])

net.load_state_dict(cpt['state_dict'])
print(net)

#%% test a trained model