EEGNet_training_LeakyReLU.py

# -*- coding: utf-8 -*-
"""
Created on Fri Jul  9 17:01:35 2021

@author: haoyuan
"""

import torch
from torch.autograd import Variable
import torch.nn.functional as F
import torch.nn as nn
from dataloader import read_bci_data
import matplotlib.pyplot as plt
from torch.utils.data import TensorDataset,DataLoader
import numpy as np
import pandas as pd
import torch.optim as optim
from torchsummary import summary
import os
import argparse

def plot_here_hustory(train_accuracy_history,test_accuracy_history,loss_history):
    
    plt.figure()

    plt.suptitle("Training Curve",fontsize=15)

    plt.subplot(1,2,1)
    plt.title("Accuracy Curve",fontsize=14)
    plt.plot(train_accuracy_history,c='r', label='Training')
    plt.plot(test_accuracy_history,c='b',label='Testing')
    plt.legend(loc='best')
    
    plt.xlabel("Epochs")

    plt.subplot(1,2,2)
    plt.title("Loss Curve",fontsize=14)
    plt.plot(loss_history,c='g')
    plt.xlabel("Epochs")

    plt.show()


def testing(x_test,y_test,model,device,filepath):

    # model.load_state_dict(torch.load(filepath))
    model.eval()
    with torch.no_grad():
        model.cuda(0)
        n = x_test.shape[0]

        x_test = x_test.astype("float32")
        y_test = y_test.astype("float32").reshape(y_test.shape[0],)
        # y_test = y_test.astype("float32").reshape(y_test.shape[0],1)

        x_test, y_test = Variable(torch.from_numpy(x_test)),Variable(torch.from_numpy(y_test))
        
        x_test,y_test = x_test.to(device),y_test.to(device)
        y_pred_test = model(x_test)

        # correct_test = (y_pred_test.ge(0.5) == y_test).sum().item()
        correct_test = (torch.max(y_pred_test,1)[1]==y_test).sum().item()
        test_accuracy = correct_test/n
        # print("testing accuracy:",correct/n)

    return test_accuracy

parser = argparse.ArgumentParser()
parser.add_argument('--epochs', type=int, default='700', help='training epochs')
parser.add_argument('--learning_rate', type=float, default='1e-3', help='learning rate')
parser.add_argument('--save_model', action='store_true', help='check if you want to save the model.')
parser.add_argument('--save_csv', action='store_true', help='check if you want to save the training history.')
opt = parser.parse_args()

torch.manual_seed(1)    # reproducible
epochs = 700
lr = 1e-3

filepath = os.path.abspath(os.path.dirname(__file__))+"\checkpoint\EEGNet_checkpoint_LeakyReLU.rar"
filepath_csv = os.path.abspath(os.path.dirname(__file__))+"\history_csv\EEGNet_LeakyReLU.csv"

min_loss=1
max_train_accuracy = 0
max_test_accuracy = 0

device = torch.device("cuda:0")

train_data, train_label, test_data, test_label = read_bci_data()

n = train_data.shape[0]

train_data = train_data.astype("float32")
train_label = train_label.astype("float32").reshape(train_label.shape[0],)
# train_label = train_label.astype("float32").reshape(train_label.shape[0],1)

# train_data.shape = (1080,1,2,750)
# train_label.shape = (1080,)

# loader = DataLoader(TensorDataset(train_data,train_label),batch_size=8)
x, y = Variable(torch.from_numpy(train_data)),Variable(torch.from_numpy(train_label))
y=torch.tensor(y, dtype=torch.long) 

class EEGNet_LeakyReLU(torch.nn.Module):
    def __init__(self, n_output):
        super(EEGNet_LeakyReLU, self).__init__()
        self.firstConv = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=(1,51), stride=(1,1), padding=(0,25),bias=False),
            nn.BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        self.depthwiseConv = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=(2,1), stride=(1,1), groups=8,bias=False),
            nn.BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),
            nn.LeakyReLU(negative_slope=0.06),
            nn.AvgPool2d(kernel_size=(1,4), stride=(1,4),padding=0),
            nn.Dropout(p=0.5)
        )
        self.separableConv = nn.Sequential(
            nn.Conv2d(32, 32, kernel_size=(1,15), stride=(1,1), padding=(0,7),bias=False),
            nn.BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),
            nn.LeakyReLU(negative_slope=0.06),
            nn.AvgPool2d(kernel_size=(1,8), stride=(1,8),padding=0),
            nn.Dropout(p=0.5)
        )
        self.classify = nn.Sequential(
            nn.Flatten(),
            nn.Linear(736,n_output,bias=True)
        )

    def forward(self, x):
        out = self.firstConv(x)
        out = self.depthwiseConv(out)
        out = self.separableConv(out) 
        out = self.classify(out)
        return out

model = EEGNet_LeakyReLU(n_output=2)
print(model)
criterion = nn.CrossEntropyLoss()

# optimizer = optim.Adam(model.parameters(),lr = lr)
optimizer = optim.RMSprop(model.parameters(),lr = lr, momentum = 0.6)
# optimizer = optim.SGD(model.parameters(), lr=1, momentum=0.5, weight_decay=5e-4)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[100,300,500], gamma=0.1)

model.cuda(0)
summary(model.cuda(),(1,2,750))

loss_history = []
train_accuracy_history = []
test_accuracy_history = []

for epoch in range(epochs):
    # for idx,(data,target) in enumerate(loader):
    model.train()
    x,y = x.to(device),y.to(device)
    y_pred = model(x)

    # print(y_pred.shape)
    # print(y.shape)

    # loss  =  F.mse_loss(y_pred, y)

    loss = criterion(y_pred, y)
    train_loss = loss.item()
    loss_history.append(train_loss)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    # scheduler.step()

    if epoch%1==0:

        # correct= (y_pred.ge(0.5) == y).sum().item()
        n = y.shape[0]
        correct = (torch.max(y_pred,1)[1]==y).sum().item()
        train_accuracy = correct / n
        train_accuracy_history.append(train_accuracy)

        # print("epochs:",epoch,"loss:",loss.item(),"Accuracy:",(correct / n),"Learning rate:",scheduler.get_last_lr()[0])
        test_accuracy = testing(test_data,test_label,model,device,filepath)
        test_accuracy_history.append(test_accuracy)

        print("epochs:",epoch,"loss:",train_loss,"Training Accuracy:",train_accuracy,"Testing Accuracy:",test_accuracy,"Learning rate:",scheduler.get_last_lr()[0])

        if train_loss<min_loss:
            min_loss = train_loss
            # torch.save(model.state_dict(), filepath)
        
        if train_accuracy>max_train_accuracy:
            max_train_accuracy = train_accuracy
            if opt.save_model:
                torch.save(model.state_dict(), filepath)
        
        if test_accuracy>max_test_accuracy:
            max_test_accuracy = test_accuracy

print("最大的Training Accuracy為:",max_train_accuracy,"最大的Testing Accuracy為:",max_test_accuracy,"最小的Loss值為:",min_loss)
plot_here_hustory(train_accuracy_history,test_accuracy_history,loss_history)
df = pd.DataFrame({"loss":loss_history,"train_accuracy_history":train_accuracy_history,"test_accuracy_history":test_accuracy_history})
if opt.save_csv:
    df.to_csv(filepath_csv,encoding="utf-8-sig")