--- a
+++ b/Models.py
@@ -0,0 +1,323 @@
+'''
+Created by Victor Delvigne
+ISIA Lab, Faculty of Engineering University of Mons, Mons (Belgium)
+victor.delvigne@umons.ac.be
+
+Source: Bashivan, et al."Learning Representations from EEG with Deep Recurrent-Convolutional Neural Networks." International conference on learning representations (2016).
+
+Copyright (C) 2019 - UMons
+
+This library is free software; you can redistribute it and/or
+modify it under the terms of the GNU Lesser General Public
+License as published by the Free Software Foundation; either
+version 2.1 of the License, or (at your option) any later version.
+
+This library is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+Lesser General Public License for more details.
+
+You should have received a copy of the GNU Lesser General Public
+License along with this library; if not, write to the Free Software
+Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
+'''
+
+import torch
+
+import torch.optim as optim
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BasicCNN(nn.Module):
+    '''
+    Build the  Mean Basic model performing a classification with CNN 
+
+    param input_image: list of EEG image [batch_size, n_window, n_channel, h, w]
+    param kernel: kernel size used for the convolutional layers
+    param stride: stride apply during the convolutions
+    param padding: padding used during the convolutions
+    param max_kernel: kernel used for the maxpooling steps
+    param n_classes: number of classes
+    return x: output of the last layers after the log softmax
+    '''
+    def __init__(self, input_image=torch.zeros(1, 3, 32, 32), kernel=(3,3), stride=1, padding=1,max_kernel=(2,2), n_classes=4):
+        super(BasicCNN, self).__init__()
+
+        n_channel = input_image.shape[1]
+
+        self.conv1 = nn.Conv2d(n_channel,32,kernel,stride=stride, padding=padding)
+        self.conv2 = nn.Conv2d(32,32,kernel,stride=stride, padding=padding)
+        self.conv3 = nn.Conv2d(32,32,kernel,stride=stride, padding=padding)
+        self.conv4 = nn.Conv2d(32,32,kernel,stride=stride, padding=padding)
+        self.pool1 = nn.MaxPool2d(max_kernel)
+        self.conv5 = nn.Conv2d(32,64,kernel,stride=stride,padding=padding)
+        self.conv6 = nn.Conv2d(64,64,kernel,stride=stride,padding=padding)
+        self.conv7 = nn.Conv2d(64,128,kernel,stride=stride,padding=padding)
+
+        self.pool = nn.MaxPool2d((1,1))
+        self.drop = nn.Dropout(p=0.5)
+
+        self.fc1 = nn.Linear(2048,512)
+        self.fc2 = nn.Linear(512,n_classes)
+        self.max = nn.LogSoftmax()
+    
+    def forward(self, x):
+        batch_size = x.shape[0]
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = F.relu(self.conv3(x))
+        x = F.relu(self.conv4(x))
+        x = self.pool1(x)
+        x = F.relu(self.conv5(x))
+        x = F.relu(self.conv6(x))
+        x = self.pool1(x)
+        x = F.relu(self.conv7(x))
+        x = self.pool1(x)
+        x = x.reshape(x.shape[0],x.shape[1], -1)
+        x = self.pool(x)
+        x = x.reshape(x.shape[0],-1)
+        x = self.fc1(x)
+        x = self.fc2(x)
+        x = self.max(x)
+        return x
+
+
+class MaxCNN(nn.Module):
+    '''
+    Build the Max-pooling model performing a maxpool over the 7 parallel convnets
+
+    param input_image: list of EEG image [batch_size, n_window, n_channel, h, w]
+    param kernel: kernel size used for the convolutional layers
+    param stride: stride apply during the convolutions
+    param padding: padding used during the convolutions
+    param max_kernel: kernel used for the maxpooling steps
+    param n_classes: number of classes
+    return x: output of the last layers after the log softmax
+    '''
+    def __init__(self, input_image=torch.zeros(1, 7, 3, 32, 32), kernel=(3,3), stride=1, padding=1,max_kernel=(2,2), n_classes=4):
+        super(MaxCNN, self).__init__()
+
+        n_window = input_image.shape[1]
+        n_channel = input_image.shape[2]
+
+        self.conv1 = nn.Conv2d(n_channel,32,kernel,stride=stride, padding=padding)
+        self.conv2 = nn.Conv2d(32,32,kernel,stride=stride, padding=padding)
+        self.conv3 = nn.Conv2d(32,32,kernel,stride=stride, padding=padding)
+        self.conv4 = nn.Conv2d(32,32,kernel,stride=stride, padding=padding)
+        self.pool1 = nn.MaxPool2d(max_kernel)
+        self.conv5 = nn.Conv2d(32,64,kernel,stride=stride,padding=padding)
+        self.conv6 = nn.Conv2d(64,64,kernel,stride=stride,padding=padding)
+        self.conv7 = nn.Conv2d(64,128,kernel,stride=stride,padding=padding)
+
+        self.pool = nn.MaxPool2d((n_window,1))
+        self.drop = nn.Dropout(p=0.5)
+
+        self.fc = nn.Linear(n_window*int(4*4*128/n_window),512)
+        self.fc2 = nn.Linear(512,n_classes)
+        self.max = nn.LogSoftmax()
+
+    def forward(self, x):
+        if x.get_device() == 0:
+            tmp = torch.zeros(x.shape[0],x.shape[1],128,4,4).cuda()
+        else:
+            tmp = torch.zeros(x.shape[0],x.shape[1],128,4,4).cpu()
+        for i in range(7):
+            tmp[:,i] = self.pool1( F.relu(self.conv7(self.pool1(F.relu(self.conv6(F.relu(self.conv5(self.pool1( F.relu(self.conv4(F.relu(self.conv3( F.relu(self.conv2(F.relu(self.conv1(x[:,i])))))))))))))))))
+        x = tmp.reshape(x.shape[0], x.shape[1],4*128*4,1)
+        x = self.pool(x)
+        x = x.view(x.shape[0],-1)
+        x = self.fc2(self.fc(x))
+        x = self.max(x)
+        return x
+
+
+class TempCNN(nn.Module):
+    '''
+    Build the Conv1D model performing a convolution1D over the 7 parallel convnets
+
+    param input_image: list of EEG image [batch_size, n_window, n_channel, h, w]
+    param kernel: kernel size used for the convolutional layers
+    param stride: stride apply during the convolutions
+    param padding: padding used during the convolutions
+    param max_kernel: kernel used for the maxpooling steps
+    param n_classes: number of classes
+    return x: output of the last layers after the log softmax
+    '''
+    def __init__(self, input_image=torch.zeros(1, 7, 3, 32, 32), kernel=(3,3), stride=1, padding=1,max_kernel=(2,2), n_classes=4):
+        super(TempCNN, self).__init__()
+
+        n_window = input_image.shape[1]
+        n_channel = input_image.shape[2]
+
+        self.conv1 = nn.Conv2d(n_channel,32,kernel,stride=stride, padding=padding)
+        self.conv2 = nn.Conv2d(32,32,kernel,stride=stride, padding=padding)
+        self.conv3 = nn.Conv2d(32,32,kernel,stride=stride, padding=padding)
+        self.conv4 = nn.Conv2d(32,32,kernel,stride=stride, padding=padding)
+        self.pool1 = nn.MaxPool2d(max_kernel)
+        self.conv5 = nn.Conv2d(32,64,kernel,stride=stride,padding=padding)
+        self.conv6 = nn.Conv2d(64,64,kernel,stride=stride,padding=padding)
+        self.conv7 = nn.Conv2d(64,128,kernel,stride=stride,padding=padding)
+
+        #Temporal CNN Layer
+        self.conv8 = nn.Conv1d(n_window,64,(4*4*128,3),stride=stride,padding=padding)
+
+        self.pool = nn.MaxPool2d((n_window,1))
+        self.drop = nn.Dropout(p=0.5)
+        self.fc = nn.Linear(64*3,n_classes)
+        self.max = nn.LogSoftmax()
+
+    def forward(self, x):
+        if x.get_device() == 0:
+            tmp = torch.zeros(x.shape[0],x.shape[1],128,4,4).cuda()
+        else:
+            tmp = torch.zeros(x.shape[0],x.shape[1],128,4,4).cpu()
+        for i in range(7):
+            tmp[:,i] = self.pool1( F.relu(self.conv7(self.pool1(F.relu(self.conv6(F.relu(self.conv5(self.pool1( F.relu(self.conv4(F.relu(self.conv3( F.relu(self.conv2(F.relu(self.conv1(x[:,i])))))))))))))))))
+        x = tmp.reshape(x.shape[0], x.shape[1],4*128*4,1)
+        x = F.relu(self.conv8(x))
+        x = x.view(x.shape[0],-1)
+        x = self.fc(x)
+        x = self.max(x)
+        return x
+
+
+class LSTM(nn.Module):
+    '''
+    Build the LSTM model applying a RNN over the 7 parallel convnets outputs
+
+    param input_image: list of EEG image [batch_size, n_window, n_channel, h, w]
+    param kernel: kernel size used for the convolutional layers
+    param stride: stride apply during the convolutions
+    param padding: padding used during the convolutions
+    param max_kernel: kernel used for the maxpooling steps
+    param n_classes: number of classes
+    param n_units: number of units
+    return x: output of the last layers after the log softmax
+    '''
+    def __init__(self, input_image=torch.zeros(1, 7, 3, 32, 32), kernel=(3,3), stride=1, padding=1,max_kernel=(2,2), n_classes=4, n_units=128):
+        super(LSTM, self).__init__()
+
+        n_window = input_image.shape[1]
+        n_channel = input_image.shape[2]
+
+        self.conv1 = nn.Conv2d(n_channel,32,kernel,stride=stride, padding=padding)
+        self.conv2 = nn.Conv2d(32,32,kernel,stride=stride, padding=padding)
+        self.conv3 = nn.Conv2d(32,32,kernel,stride=stride, padding=padding)
+        self.conv4 = nn.Conv2d(32,32,kernel,stride=stride, padding=padding)
+        self.pool1 = nn.MaxPool2d(max_kernel)
+        self.conv5 = nn.Conv2d(32,64,kernel,stride=stride,padding=padding)
+        self.conv6 = nn.Conv2d(64,64,kernel,stride=stride,padding=padding)
+        self.conv7 = nn.Conv2d(64,128,kernel,stride=stride,padding=padding)
+
+        # LSTM Layer
+        self.rnn = nn.RNN(4*4*128, n_units, n_window)
+        self.rnn_out = torch.zeros(2, 7, 128)
+
+        self.pool = nn.MaxPool2d((n_window,1))
+        self.drop = nn.Dropout(p=0.5)
+        self.fc = nn.Linear(896, n_classes)
+        self.max = nn.LogSoftmax()
+
+    def forward(self, x):
+        if x.get_device() == 0:
+            tmp = torch.zeros(x.shape[0], x.shape[1], 128, 4, 4).cuda()
+        else:
+            tmp = torch.zeros(x.shape[0], x.shape[1], 128, 4, 4).cpu()
+        for i in range(7):
+            img = x[:, i]
+            img = F.relu(self.conv1(img))
+            img = F.relu(self.conv2(img))
+            img = F.relu(self.conv3(img))
+            img = F.relu(self.conv4(img))
+            img = self.pool1(img)
+            img = F.relu(self.conv5(img))
+            img = F.relu(self.conv6(img))
+            img = self.pool1(img)
+            img = F.relu(self.conv7(img))
+            tmp[:, i] = self.pool1(img)
+            del img
+        x = tmp.reshape(x.shape[0], x.shape[1], 4 * 128 * 4)
+        del tmp
+        self.rnn_out, _ = self.rnn(x)
+        x = self.rnn_out.view(x.shape[0], -1)
+        x = self.fc(x)
+        x = self.max(x)
+        return x
+
+
+class Mix(nn.Module):
+    '''
+        Build the LSTM model applying a RNN and a CNN over the 7 parallel convnets outputs
+
+        param input_image: list of EEG image [batch_size, n_window, n_channel, h, w]
+        param kernel: kernel size used for the convolutional layers
+        param stride: stride apply during the convolutions
+        param padding: padding used during the convolutions
+        param max_kernel: kernel used for the maxpooling steps
+        param n_classes: number of classes
+        param n_units: number of units
+        return x: output of the last layers after the log softmax
+        '''
+    def __init__(self, input_image=torch.zeros(1, 7, 3, 32, 32), kernel=(3,3), stride=1, padding=1,max_kernel=(2,2), n_classes=4, n_units=128):
+        super(Mix, self).__init__()
+
+        n_window = input_image.shape[1]
+        n_channel = input_image.shape[2]
+
+        self.conv1 = nn.Conv2d(n_channel,32,kernel,stride=stride, padding=padding)
+        self.conv2 = nn.Conv2d(32,32,kernel,stride=stride, padding=padding)
+        self.conv3 = nn.Conv2d(32,32,kernel,stride=stride, padding=padding)
+        self.conv4 = nn.Conv2d(32,32,kernel,stride=stride, padding=padding)
+        self.pool1 = nn.MaxPool2d(max_kernel)
+        self.conv5 = nn.Conv2d(32,64,kernel,stride=stride,padding=padding)
+        self.conv6 = nn.Conv2d(64,64,kernel,stride=stride,padding=padding)
+        self.conv7 = nn.Conv2d(64,128,kernel,stride=stride,padding=padding)
+
+        # LSTM Layer
+        self.rnn = nn.RNN(4*4*128, n_units, n_window)
+        self.rnn_out = torch.zeros(2, 7, 128)
+
+        # Temporal CNN Layer
+        self.conv8 = nn.Conv1d(n_window, 64, (4 * 4 * 128, 3), stride=stride, padding=padding)
+
+        self.pool = nn.MaxPool2d((n_window, 1))
+        self.drop = nn.Dropout(p=0.5)
+        self.fc1 = nn.Linear(1088,512)
+        self.fc2 = nn.Linear(512, n_classes)
+        self.max = nn.LogSoftmax()
+
+
+    def forward(self, x):
+        if x.get_device() == 0:
+            tmp = torch.zeros(x.shape[0], x.shape[1], 128, 4, 4).cuda()
+        else:
+            tmp = torch.zeros(x.shape[0], x.shape[1], 128, 4, 4).cpu()
+        for i in range(7):
+            img = x[:, i]
+            img = F.relu(self.conv1(img))
+            img = F.relu(self.conv2(img))
+            img = F.relu(self.conv3(img))
+            img = F.relu(self.conv4(img))
+            img = self.pool1(img)
+            img = F.relu(self.conv5(img))
+            img = F.relu(self.conv6(img))
+            img = self.pool1(img)
+            img = F.relu(self.conv7(img))
+            tmp[:, i] = self.pool1(img)
+            del img
+
+        temp_conv = F.relu(self.conv8(tmp.reshape(x.shape[0], x.shape[1], 4 * 128 * 4, 1)))
+        temp_conv = temp_conv.reshape(temp_conv.shape[0], -1)
+
+        self.lstm_out, _ = self.rnn(tmp.reshape(x.shape[0], x.shape[1], 4 * 128 * 4))
+        del tmp
+        lstm = self.lstm_out.view(x.shape[0], -1)
+
+        x = torch.cat((temp_conv, lstm), 1)
+
+        x = self.fc1(x)
+        x = self.fc2(x)
+        x = self.max(x)
+        return x
+