EEGNet复现

一、前言

本篇文章记录一下本人复现EEGNet的记录，文章包含三个文件（模块），第一个是脑电数据的预处理，EEGNet模型构建，EEGNet训练模块。

本次实验的数据集是BCIcompetitionIVdata2a数据集

应用的库：mne torch scipy numpy sklearn

二、正文

1.数据预处理

先展示一下编写的数据预处理文件

import mne
import numpy as np
import torch
from sklearn.preprocessing import StandardScaler,OneHotEncoder
import scipy.io
import torchvision.transforms as transforms

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

"""进行预处理前请检查数据是否存在过多的噪声，或者数据是否稳定"""

# 进行数据预处理 只适合T结尾的data数据，其他的需要修改函数 filename需以.npy结尾
def transform_save_data(filename,save_filename):
    raw = mne.io.read_raw_gdf(filename)
    print(raw.info['ch_names'])
    events,event_id = mne.events_from_annotations(raw)
    raw.info['bads'] += ['EOG-left','EOG-central','EOG-right']
    # 运动想象时间2-6秒
    tmin,tmax = 2,6 
    event_id = {'769':7,'770':8,'771':9,'772':10}
    # 需要重新加载raw对象进行滤波处理
    raw.load_data()
    raw.filter(7.0,35.0,fir_design='firwin')
    picks = mne.pick_types(raw.info,meg=False,eeg=True,stim=False,exclude='bads')
    epochs = mne.Epochs(raw=raw,events=events,event_id=event_id,tmin=tmin,tmax=tmax,preload=True,baseline=None,picks=picks)
    epoch_data = epochs.get_data()
    # 将最后一位数据进行去除
    epoch_data = epoch_data[:,:,:-1]


    epoch_data = epoch_data.reshape(epoch_data.shape[0], 1, 22, 1000)
    np.save(save_filename,epoch_data)


# 进行归一化处理
def data_processing(BCI_IV_2a_data,label_filename):
    Scaler = StandardScaler()
    X_train = BCI_IV_2a_data.reshape(BCI_IV_2a_data.shape[0], 22000)
    X_train_Scaler = Scaler.fit_transform(X_train)
    # 进行reshape第二个维度为channels W H
    acc_train = X_train_Scaler.reshape(BCI_IV_2a_data.shape[0], 1, 22, 1000)
    data_label = scipy.io.loadmat(label_filename)
    Label = data_label['classlabel']
    # 查看由多少类
    n_classes = len(np.unique(Label))
    print(n_classes)
    encoder = OneHotEncoder(handle_unknown='ignore')
    y = np.array(Label)
    y_oh = encoder.fit_transform(y).toarray()
    y_oh = y_oh[:-1]

    return acc_train,y_oh


# 进行转换成Tensor格式的数据  保存的文件的格式应该以pt为后缀
def data_transform_tensor(acc_train,y_oh,save_datafilename,save_labelfilename):
    transf = transforms.ToTensor()
    d = transf(y_oh)
    # 去除另外四个维度的标签，标签就是最大值
    label = torch.argmax(d,dim=2).long()

    

    h = torch.squeeze(label)
    data = torch.tensor(acc_train,dtype=torch.float32)
    labels = torch.tensor(h,dtype=torch.long)

    torch.save(data,save_datafilename)
    torch.save(labels,save_labelfilename)



# 将数据进行联合
def combine_data(data_list,label_list,data_filename,label_filename):
    """_summary_
    将增强的EEG_data数据进行拼接 并保存为pt后缀文件
    Parameters
    ----------
    data_list : _type_  tensor
        EEGdata list
    label_list : _type_ tensor
        label list
    data_filename : _type_, optional
        _description_, by default None 
    label_filename : _type_, optional
        _description_, by default None 
    """
    data_combine = torch.cat(data_list, axis=0)
    label_combine = torch.cat(label_list, axis=0)
    torch.save(data_combine, data_filename)
    torch.save(label_combine, label_filename)




# 进行时域上EEG数据增强  通过分割，重构 打乱数据

def interaug(timg, label, batch_size):
    """timg是data label是标签"""
    aug_data = []
    aug_label = []
    for cls4aug in range(4):
        # 条件判断 找出对应的label和data
        cls_idx = np.where(label == cls4aug)  # label == cls4aug + 1
        tmp_data = timg[cls_idx]
        tmp_label = label[cls_idx]
        # 分epoch
        tmp_aug_data = np.zeros((int(batch_size / 4), 1, 22, 1000))
        for ri in range(int(batch_size / 4)):
            # 随机取8个时间片段
            for rj in range(8):
                rand_idx = np.random.randint(0, tmp_data.shape[0], 8)
                # 进行数据的打乱重构
                tmp_aug_data[ri, :, :, rj * 125:(rj + 1) * 125] = tmp_data[rand_idx[rj], :, :, rj * 125:(rj + 1) * 125]

        aug_data.append(tmp_aug_data)
        aug_label.append(tmp_label[:int(batch_size / 4)])
    aug_data = np.concatenate(aug_data)
    aug_label = np.concatenate(aug_label)
    aug_shuffle = np.random.permutation(len(aug_data))
    aug_data = aug_data[aug_shuffle, :, :]
    aug_label = aug_label[aug_shuffle]

    aug_data = torch.from_numpy(aug_data).cuda()
    aug_data = aug_data.float()
    aug_label = torch.from_numpy(aug_label).cuda()  # aug_label - 1
    aug_label = aug_label.long()
    return aug_data, aug_label



# 切分部分数据进行test
def split_EEGdata(data, label):
    data = data.view(data.shape[0], 1, 22, 1000)
    data = data[:100]
    label = label[:100]
    torch.save(data, 'EEG_data_split.pt')
    torch.save(label, 'EEG_label_split.pt')



def transform_save_data_version2(filename,save_filename):
    raw = mne.io.read_raw_gdf(filename)
    print(raw.info['ch_names'])
    events,event_id = mne.events_from_annotations(raw)
    raw.info['bads'] += ['EOG-left','EOG-central','EOG-right']
    # 运动想象时间2-6秒
    tmin,tmax = 2,6
    event_id = {'769':7,'770':8,'771':9,'772':10}
    # 需要重新加载raw对象进行滤波处理
    raw.load_data()
    picks = mne.pick_types(raw.info,meg=False,eeg=True,stim=False,exclude='bads')
    epochs = mne.Epochs(raw=raw,events=events,event_id=event_id,tmin=tmin,tmax=tmax,preload=True,baseline=None,picks=picks)
    epoch_data = epochs.get_data()
    # 将最后一位数据进行去除
    epoch_data = epoch_data[:,:,:-1]


    epoch_data = epoch_data.reshape(epoch_data.shape[0], 1, 22, 1000)
    np.save(save_filename,epoch_data)



if __name__ == '__main__':
    # filename = 'C:\\Users\\24242\\Desktop\\AI_Reference\\data_bag\\BCICIV_2a_gdf\\A06T.gdf'
    # save_filename = 'A06T.npy'
    # transform_save_data(filename,save_filename)
    # data = np.load('C:\\Users\\24242\\DataspellProjects\\EEG_Project\\EEGNet\\A06T.npy')
    # label_filename = 'C:\\Users\\24242\\Desktop\\AI_Reference\\data_bag\\BCICIV_2a_gdf\\A06T.mat'
    # acc_train,y_oh = data_processing(data,label_filename)
    # save_datafilename = 'A06T.pt'
    # save_labelfilename = 'A06T_target.pt'
    # data_transform_tensor(acc_train,y_oh,save_datafilename,save_labelfilename)
    # 进行数据的拼接 使深度学习模型训练样本不会太少
     
    # data1 = np.load('C:\\Users\\24242\\DataspellProjects\\EEG_Project\\EEGNet\\A01T.npy')
    # data2 = torch.load('C:\\Users\\24242\\DataspellProjects\\EEG_Project\\EEGNet\\A02T.pt')
    # data3 = torch.load('C:\\Users\\24242\\DataspellProjects\\EEG_Project\\EEGNet\\A03T.pt')
    # data5 = torch.load('C:\\Users\\24242\\DataspellProjects\\EEG_Project\\EEGNet\\A05T.pt')
    # label1 = torch.load('C:\\Users\\24242\\DataspellProjects\\EEG_Project\EEGNet\\A01T_target.pt')
    # label1 = label1[:-1]
    # label2 = torch.load('C:\\Users\\24242\\DataspellProjects\\EEG_Project\EEGNet\\A02T_target.pt')
    # label2 = label2[:-1]
    # label3 = torch.load('C:\\Users\\24242\\DataspellProjects\\EEG_Project\EEGNet\\A03T_target.pt')
    # label3 = label3[:-1]
    # label1 = label1.numpy()
    # label5 = torch.load('C:\\Users\\24242\\DataspellProjects\\EEG_Project\EEGNet\\A05T_target.pt')
    
    # data_list = [data1,data2,data3,data5]
    # label_list = [label1,label2,label3,label5]
    # data_filename = 'combine_data_01.pt'
    # label_filename = 'combine_label_01.pt'
    # label1 = label1[:-1]
    #
    # aug_data, aug_label = interaug(timg=data1,label=label1,batch_size=287)
    #
    # data_1 = torch.load('C:\\Users\\24242\\DataspellProjects\\EEG_Project\EEGNet\\A01T.pt')
    # label_1 = torch.load('C:\\Users\\24242\\DataspellProjects\\EEG_Project\EEGNet\\A01T_target.pt')
    # print(label_1.shape)
    # print(data_1.shape)
    # print(aug_data.shape)
    # print(aug_label.shape)
    # data_list = [data_1,aug_data]
    # label_list = [label_1,aug_label]
    # combine_data(data_list=data_list,label_list=label_list,data_filename=data_filename,label_filename=label_filename)
    print('------------------new FileProcess------------------')
    # filename = 'C:\\Users\\24242\\Desktop\\AI_Reference\\data_bag\\BCICIV_2a_gdf\\A01T.gdf'
    # save_filename = 'A01T_new_original.npy'
    # # transform_save_data_version2(filename, save_filename)
    # data = np.load(save_filename)
    #
    # label_name = 'C:\\Users\\24242\\Desktop\\AI_Reference\\data_bag\\BCICIV_2a_gdf\\A01T.mat'
    # acc_train, y_oh = data_processing(data,label_name)
    #
    # save_dataFilename = 'A01T_new_original.pt'
    # save_labelFilename = 'A01T_new_label_original.pt'
    # # data_transform_tensor(acc_train, y_oh, save_dataFilename, save_labelFilename)
    data_final = torch.load('A01T_new_original.pt').to('cpu')
    label_final = torch.load('A01T_new_label_original.pt').to('cpu')
    # data_final_01 = torch.load('A01T_new_data_aug.pt').to('cpu')
    # label_final_01 = torch.load('A01T_new_label_aug.pt').to('cpu')
    data_final_02 = torch.load('A01T_new_data_aug_01.pt').to('cpu')
    label_final_02 = torch.load('A01T_new_label_aug_01.pt').to('cpu')
    data_final_03 = torch.load('A01T_new_data_aug_02.pt').to('cpu')
    label_final_03 = torch.load('A01T_new_label_aug_02.pt').to('cpu')

    # data_final = data_final.numpy().astype('float32')
    # label_final = label_final.numpy().astype('float32')
    # print(label_final)
    # split_EEGdata(data_final, label_final)
    # aug_data, aug_label = interaug(data_final, label_final, batch_size=48)
    # print(aug_label)
    # torch.save(aug_data, 'A01T_new_data_aug_03.pt')
    # torch.save(aug_label, 'A01T_new_label_aug_03.pt')
    data_list = [data_final_02, data_final_03, data_final]
    label_list = [label_final_02, label_final_03, label_final]
    data_combine_filename = 'A01T_new_combine.pt'
    data_combine_label_filename = 'A01T_new_label_combine.pt'
    combine_data(data_list, label_list, data_combine_filename, data_combine_label_filename)
    label = torch.load('A01T_new_label_combine.pt')
    data = torch.load('A01T_new_combine.pt')
    # print(label)

大部分的代码都编写了注释，该模块可以直接进行应用，数据经过预处理后。

保存shape为(288,1,22,1000)的数据到文件中进行备用 

数据shape的原始格式为(trial,1,channel,sample)

2.模型构建

先上代码

# 导入工具包
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary


class EEGNet(nn.Module):
    def __init__(self, classes_num):
        super(EEGNet, self).__init__()
        self.drop_out = 0.25

        self.block_1 = nn.Sequential(
            # Pads the input tensor boundaries with zero
            # left, right, up, bottom
            nn.ZeroPad2d((8, 8, 0, 0)),
            nn.Conv2d(
                in_channels=1,  # input shape (1, C, T)
                out_channels=8,  # num_filters
                kernel_size=(1, 16),  # filter size 8HZ以上的滤波用于运动想象
                bias=False
            ),  # output shape (8, C, T)
            nn.BatchNorm2d(8)  # output shape (8, C, T)
        )

        # block 2 and 3 are implementations of Depthwise Convolution and Separable Convolution
        self.block_2 = nn.Sequential(
            nn.Conv2d(
                in_channels=8,  # input shape (8, C, T)
                out_channels=16,  # num_filters
                kernel_size=(22, 1),  # filter size
                # group8意味着八组滤波器
                groups=8,
                bias=False
            ),  # output shape (16, 1, T)
            nn.BatchNorm2d(16),  # output shape (16, 1, T)
            nn.ELU(),
            nn.AvgPool2d((1, 4)),  # output shape (16, 1, T//4)
            nn.Dropout(self.drop_out)  # output shape (16, 1, T//4)
        )

        self.block_3 = nn.Sequential(
            nn.ZeroPad2d((8, 8, 0, 0)),
            nn.Conv2d(
                in_channels=16,  # input shape (16, 1, T//4)
                out_channels=16,  # num_filters
                kernel_size=(1, 16),  # filter size
                # 十六组滤波器
                groups=16,
                bias=False
            ),  # output shape (16, 1, T//4)
            nn.Conv2d(
                in_channels=16,  # input shape (16, 1, T//4)
                out_channels=16,  # num_filters
                kernel_size=(1, 1),  # filter size
                bias=False
            ),  # output shape (16, 1, T//4)
            nn.BatchNorm2d(16),  # output shape (16, 1, T//4)
            nn.ELU(),
            nn.AvgPool2d((1, 8)),  # output shape (16, 1, T//32)
            nn.Dropout(self.drop_out)
        )

        self.out = nn.Linear((16 * 31), classes_num)

    def forward(self, x):
        x = self.block_1(x)
        # print("block1", x.shape)
        x = self.block_2(x)
        # print("block2", x.shape)
        x = self.block_3(x)
        # print("block3", x.shape)

        x = x.view(x.size(0), -1)
        x = self.out(x)
        # return F.softmax(x, dim=1), x  # return x for visualization
        return x
if __name__ == '__main__':
    input = torch.randn(32, 1, 22, 1000).to(device='cuda')

    model = EEGNet(4).to(device='cuda')
    summary(model, input_size=(1, 22, 1000), batch_size=32, device='cuda')

下面是论文中的图片

本次模型生成了8个时间滤波器，对于每个时间滤波器生成2个空间滤波器即16个空间滤波器

原论文提到dropout设置为0.25效果较好，且另外一种结构是生成4个时间滤波器，对于每个时间滤波器生成2个空间滤波器，大家可以自行测试

3.模型训练

import torch.nn
import torchvision.transforms as transforms
from torch import optim
from torch.utils.data import DataLoader
from EEGNet_pytorch_version import *
import pandas as pd

 

# 加载EEGnetdata
EEGnetdata = EEGNetDataset(file_path ='C:\\Users\\24242\\DataspellProjects\\EEG_Project\\EEGNet\\A03T.pt',
                            target_path ='C:\\Users\\24242\\DataspellProjects\\EEG_Project\\EEGNet\\A03T_target.pt',
                            transform=False,target_transform=False)


train_dataloader = DataLoader(EEGnetdata,shuffle=False,num_workers=0,batch_size=Config.train_batch_size,drop_last=True)

# 使用cuda的GPU去跑
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

print(device)
# 构建CNNnet
net = EEGNet().to(device)

# 损失函数
# criterion = torch.nn.MSELoss(reduction='mean').to(device)
criterion = torch.nn.CrossEntropyLoss().to(device)
# criterion = nn.MultiMarginLoss()
# optimizer = optim.SGD(net.parameters(), lr=0.8)
# 优化器
optimizer = optim.Adam(net.parameters(), lr=0.001)
counter = []
# 损失函数记录
loss_history = []
# 迭代次数
iteration_number = 0
train_correct = 0
total = 0
# 训练精度记录
train_accuracy = []
correct = 0
total = 0
# 分类数
classnum = 4
# 精确度记录
accuracy_history = []
# 继承了Model的train函数进行训练
# 当调用 train() 函数时，模型中的所有具有学习参数的层（例如卷积层、全连接层等）将被设置为训练模式。这意味着在模型的前向传播过程中，这些层将考虑到梯度的计算，并且在反向传播时会更新它们的权重和偏置。
net.train()
# 本次代码训练的是100个次
for epoch in range(0, Config.train_number_epochs):
    # index与element
    for i,data in enumerate(train_dataloader,0): #enumerate防止重复抽取到相同数据，数据取完就可以结束一个epoch
        # 从对象data中取到EEG_data和target
        item,target = data

        # 是否用GPU进行训练
        # item(48,22,20,50) 第一个维度epoch 第二个维度channels 第三个第四个维度为H*W
        item,target= item.to(device),target.to(device)
        print(item.shape)
        optimizer.zero_grad() #grad归零

        output = net(item)  #输出

        loss = criterion(output,target) #算loss 为损失函数,target原先为Tensor类型，指定target为long类型即可。
        loss.backward()   #反向传播算当前grad

        torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=1.0)
        optimizer.step()  #optimizer更新参数
        #求ACC标准流程
        predicted=torch.argmax(output, 1)
        train_correct += (predicted == target).sum().item()
        total+=target.size(0) # total += target.size
        train_accuracy = train_correct / total
        train_accuracy = np.array(train_accuracy)
        
        if i % 10 == 0: #每10个epoch输出一次结果
                print("Epoch number {}\n Current Accuracy {}\n Current loss {}\n".format
                      (epoch, train_accuracy.item(),loss.item()))
        iteration_number += 1
        counter.append(iteration_number)
        accuracy_history.append(train_accuracy.item())
        loss_history.append(loss.item())
        
show_plot(counter,accuracy_history,loss_history)  
        



# 保存模型 参数
# torch.save(net.state_dict(),"The train_new_EEGNet.ph")
# 保存模型
# torch.save(net, "The train_new_EEGNet_A03.ph")

三、总结

闲下来水一篇文章，本篇文章只是一个小demo，是本人模型的复现记录，并没有进行测试，大家可以跑一跑模型，继续改进一下