Commit 66c2265c authored by Ard Kastrati's avatar Ard Kastrati
Browse files

Created a sample CNN.

parent c9da4d2e
......@@ -3,9 +3,7 @@ from config import CNN_params as CNN_params
import numpy as np
import scipy.io as sio
import os
# import torch as torch
train_data_path = params['train_data_path']
import torch
def get_data(verbose=True):
"""
......@@ -14,13 +12,15 @@ def get_data(verbose=True):
:param variable: variable of the matlab file
:return: the data as numpy array / Tensor?
"""
train_data = torch_training_data(train_data_path, verbose)
return train_data
train_x = torch_data(data_path=params['data_path'], filename=CNN_params['trainX_filename'], variable1=CNN_params['trainX_variable1'], variable2=CNN_params['trainX_variable2'], verbose=verbose, detailed_verbose=True)
train_y = torch_data(data_path=params['data_path'], filename=CNN_params['trainY_filename'], variable1=CNN_params['trainY_variable1'], variable2=CNN_params['trainY_variable2'], verbose=verbose, detailed_verbose=True)
train_y = train_y.type(torch.LongTensor)
return train_x, train_y
def torch_training_data(train_data_path, verbose=True):
def torch_data(data_path, filename, variable1, variable2, verbose=True, detailed_verbose=False):
"""
Extract training data from the file.
:param train_data_path: name of the file to open.
Extract data from the file.
:param data_path: name of the file to open.
:param verbose: boolean; if true, it prints information about
the status of the program.
:return: a numpy array of shape ...?
......@@ -30,18 +30,23 @@ def torch_training_data(train_data_path, verbose=True):
trials = extract_trials()
if verbose: print(len(trials), " trials found.")
full_data = []
for i in range(5):
next_trial = load_matlab_trial(datapath=train_data_path, trial=trials[i], filename=CNN_params['filename'], variable=CNN_params['data_variable'])
full_data.append(next_trial.tolist()[0][0])
full_data = np.array([])
for i in range(len(trials)):
if detailed_verbose: print("Trying trial", trials[i])
next_trial = load_matlab_trial(datapath=data_path, trial=trials[i], filename=filename, variable1=variable1, variable2=variable2)
if full_data.size == 0:
full_data = next_trial
else:
full_data = np.concatenate((full_data, next_trial))
if verbose: print(np.shape(full_data))
if verbose: print("data loaded.")
print(np.shape(full_data))
if verbose: print("Tensoring data..")
if verbose: print("not implemented yet")
return full_data
if verbose: print("Tensoring data...")
full_data_tensor = torch.from_numpy(full_data)
if verbose: print("Data in tensor form.")
return full_data_tensor
def load_matlab_trial(datapath, trial, filename, variable):
def load_matlab_trial(datapath, trial, filename, variable1, variable2):
"""
Load the data from Matlab
:param datapath: matlab data directory
......@@ -49,13 +54,18 @@ def load_matlab_trial(datapath, trial, filename, variable):
:param variable: variable of the matlab file
:return: the data as numpy array
"""
data = sio.loadmat(datapath + trial + "/" + filename)[filename][variable]
data = sio.loadmat(datapath + trial + "/" + filename)[variable1][variable2][0][0]
if len(np.shape(data)) == 3:
data = np.swapaxes(data, 0, 2)
data = np.swapaxes(data, 1, 2)
else:
data = data - 1 # data needs to be between 0 and 1
return data
def extract_trials():
"""
Extracts the trials from the root directory
"""
my_list = os.listdir(train_data_path)
my_list = os.listdir(params['data_path'])
trials = [name for name in my_list if len(name) == 3]
return trials
\ No newline at end of file
return trials
......@@ -12,8 +12,13 @@ RNN? Transformers, Attention, etc etc
# general parameters
general_params['model'] = "CNN-1"
general_params['train_data_path'] = "/Volumes/methlab/ETH_AS/preprocessed2/"
general_params['data_path'] = "/Volumes/methlab/ETH_AS/preprocessed2/"
# CNN - 1
CNN_params['data_variable'] = "data"
CNN_params['filename'] = "EEGprocue"
CNN_params['trainX_variable1'] = "EEGprocue"
CNN_params['trainX_variable2'] = "data"
CNN_params['trainX_filename'] = "EEGprocue"
CNN_params['trainY_variable1'] = "trialinfopro"
CNN_params['trainY_variable2'] = "cues"
CNN_params['trainY_filename'] = "trialinfocuelocked"
\ No newline at end of file
from config import general_params as params
from IOHelper import get_data
import time
from run import run
def main():
start_time = time.time()
if params['model'] == 'CNN-1':
print("Started running CNN-1. If you want to run other methods please choose another model in the config.py file.")
get_data()
run()
else:
raise Exception('Please choose one of the following models in the config.py file')
......
# -*- coding: utf-8 -*-
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from IOHelper import get_data
import torch.utils.data
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv1d(in_channels=129, out_channels=258, kernel_size=5, stride=1, padding=2)
self.pool = nn.MaxPool1d(5)
self.conv2 = nn.Conv1d(in_channels=258, out_channels=64, kernel_size=5, stride=1, padding=2)
self.fc1 = nn.Linear(64*20, 120)
self.fc2 = nn.Linear(120, 60)
self.fc3 = nn.Linear(60, 2)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 64*20)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
def run(verbose=True):
#load the data
trainX, trainY = get_data()
print(trainX.shape, trainY.shape)
dataset = torch.utils.data.TensorDataset(trainX, trainY)
trainloader = torch.utils.data.DataLoader(dataset, batch_size=2)
# define the network
net = Net()
# define the optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# train
train(trainloader=trainloader, net=net, optimizer=optimizer, criterion=criterion)
# save our trained model
PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH)
def train(trainloader, net, optimizer, criterion, epoch=50):
for epoch in range(2): # loop over the dataset multiple times
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
outputs = net(inputs)
loss = criterion(outputs, labels.squeeze(1))
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if i % 200 == 0: # print every 2000 mini-batches
print('[%d, %5d] loss: %.3f' %
(epoch + 1, i + 1, running_loss / 2000))
running_loss = 0.0
print('Finished Training')
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment