EEG深度学习实战：脑电信号通过CNN神经网络进行二分类_cnn 脑电数据

2023/3/11 -3/13 脑机接口学习内容一览：

        

        这一篇文章主要建立在前文处理EEG脑电信号生成时频图dataset的基础上对数据集进行操作，通过CNN神经网络做到二分类，尝试使用深度学习的模式来预测分类脑电信号。

---

一、需求分析

        我们需要在前文生成的数据集的基础上建立自己的数据集，并稍微修改参考资料【3】中提供的神经网络，用我们的训练集来修正网络参数，最后通过测试集来检验准确率。

1.数据集

        数据集结构如图：

 

         1代表rt事件，2代表square事件，具体参考前文。

2.输出数据

        测试集的预测准确率。

---

二、处理流程

        1.导入数据集；

        2.根据数据集结构构建自己的dataset class；

        3.构建神经网络，导入训练集进行训练，并使用tensorboard查看模型结构；

        4.导入测试集，输出准确率。

---

三、完整代码

import torch.optim
from torch.utils.data import Dataset
from torchvision import transforms
from torch.utils.data import DataLoader
from PIL import Image
import os
import torch.nn as nn
from torch.utils.tensorboard import SummaryWriter
 
 
class tf_dataset(Dataset):
    def __init__(self, data_dir, transform=None):
        """
        分类任务的Dataset
        :param data_dir: str, 数据集所在路径
        :param transform: torch.transform，数据预处理
        """
        self.label_name = {"1": 0, "2": 1}
        self.data_info = self.get_img_info(data_dir)  # data_info存储所有图片路径和标签，在DataLoader中通过index读取样本
        self.transform = transform
 
    def __getitem__(self, index):
        path_img, label = self.data_info[index]
        img = Image.open(path_img).convert('RGB')     # 0~255
 
        if self.transform is not None:
            img = self.transform(img)   # 在这里做transform，转为tensor等等
        return img, label
 
    def __len__(self):
        # 返回数据长度（图片数量）
        return len(self.data_info)
 
    @staticmethod
    def get_img_info(data_dir):
        data_info = list()
        # 子文件夹（类别）分别对应的类别标签
        label_name = {"1": 0, "2": 1}
        for root, dirs, _ in os.walk(data_dir):
            # 深度遍历类别，本质上是遍历root文件夹下所有子文件
            for sub_dir in dirs:
                # 将子文件夹下的所有图片文件名称str列入list中
                img_names = os.listdir(os.path.join(root, sub_dir))
                # 筛选所有png类型图片
                img_names = list(filter(lambda x: x.endswith('.png'), img_names))
                # 遍历图片list
                for i in range(len(img_names)):
                    img_name = img_names[i]
                    # 获取当前图片路径
                    path_img = os.path.join(root, sub_dir, img_name)
                    # 根据图片的父文件夹的名称获取图片所属类别
                    label = label_name[sub_dir]
                    # 增加（图片路径, 类别）元组
                    data_info.append((path_img, label))
 
        return data_info
 
 
# 构造神经网络
class my_cnn(nn.Module):
    def __init__(self):
        super(my_cnn, self).__init__()
        self.module = nn.Sequential(
            nn.Conv2d(3, 3, 5, 1, 2),
            nn.MaxPool2d(5),
            nn.Conv2d(3, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64*4*4, 64),
            nn.Linear(64, 2)
        )
 
    def forward(self, x):
        x = self.module(x)
        return x
 
 
# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device = "cpu"
cnn_model = my_cnn().to(device)
print(cnn_model)
 
"""
1.导入数据集
"""
split_dir = os.path.join(os.getcwd(), 'data/root')
train_dir = os.path.join(split_dir, 'train')
test_dir = os.path.join(split_dir, 'test')
 
# 构建transforms预处理方式
trans_resize = transforms.CenterCrop((160, 160))
trans_totensor = transforms.ToTensor()
trans_normal = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
pre_transform = transforms.Compose([trans_resize, trans_totensor, trans_normal])
 
# 创建可用数据集
train_data = tf_dataset(data_dir=train_dir, transform=pre_transform)
test_data = tf_dataset(data_dir=test_dir, transform=pre_transform)
print("训练集长度为：{}".format(len(train_data)))
print("测试集长度为：{}".format(len(test_data)))
 
# 构建DataLoader
batch_size = 32     # 确定batch_size
train_loader = DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True)
print(train_data[0][0].size())
test_loader = DataLoader(dataset=test_data, batch_size=batch_size)
 
# 确定损失函数
loss_fn = nn.CrossEntropyLoss()
 
# 创建网络模型
EEGnet = my_cnn()
 
# 优化器
learning_rate = 0.008
optimizer = torch.optim.SGD(EEGnet.parameters(), lr=learning_rate)
 
# 训练轮数
epoch = 50
# 已训练轮数
total_train_step = 0
# 已测试轮数
total_test_step = 0
# 添加tensorboard
writer = SummaryWriter("./logs_train")
 
for i in range(epoch):
    print("----------第{}轮训练开始----------".format(i+1))
    EEGnet.train()
    for data in train_loader:
        imgs, targets = data
        # print(imgs.size())
        # print(targets)
        # 数据导入模型生成结果
        outputs = EEGnet(imgs)
        # print(outputs.shape)
        loss = loss_fn(outputs, targets)
 
        # 优化器优化模型
        optimizer.zero_grad()   # 清除优化梯度
        loss.backward()         # 反向传播计算梯度值
        optimizer.step()        # 参数更新
 
        total_train_step += 1
        # 训练50次打印1次
        if total_train_step % 10 == 0:
            print("训练次数：{}, loss = {}".format(total_train_step, loss.item()))
            writer.add_scalar("train_loss", loss.item(), total_train_step)
 
    # 测试（验证）步骤开始(评估模型是否训练好)
    EEGnet.eval()
    total_test_loss = 0
    total_accuracy = 0
    # 若梯度近似为0
    with torch.no_grad():
        for data in test_loader:
            imgs, targets = data
            outputs = EEGnet(imgs)
            loss = loss_fn(outputs, targets)
            total_test_loss += loss.item()
            # 求取准确率
            accuracy = (outputs.argmax(1) == targets).sum()
            total_accuracy += accuracy
    print("整体测试集上的Loss = {}".format(total_test_loss))
    print("整体测试集上的正确率 = {}".format(total_accuracy*1.0/len(test_data)))
 
    writer.add_scalar("test_loss", total_test_loss, total_test_step)
    writer.add_scalar("test_accuracy", total_accuracy*1.0/len(test_data), total_test_step)
    total_test_step += 1
 
    # 保存模型参数,以训练轮数区分
    torch.save(EEGnet, "model_save/EEGnet_{}.path".format(i))
    print("模型已保存")
 
 
writer.close()

---

四、结果展示与分析 

        根据参考资料【4】中对于命令行中tensorboard的操作方法，我们可以在开启的tensorboard界面中获取这样的训练过程图像：

 

 

        从这里可以看出训练的误差在不断减小。 

        

        50轮训练稍微有点多了，容易过拟合，最后出来的准确率其实不太好看。在思考之后总结出了以下几点可能的错误：
        将通道分离开绘制同一epoch下的各个通道的时频图，然后进行训练预测有一些不妥当，在同一类别（rt, square）下的各通道的时频图像很显然属于不同的分布，在之后会尝试一下使用平均诱发电位evoke进行绘制和预测。

---

五、参考文献

【1】Pytorch数据读取机制(DataLoader)与图像预处理模块(transforms)

【2】CNN数据集——自己建立数据集要点

【3】脑电EEG代码开源分享 【6. 分类模型-深度学习篇】

【4】PyTorch深度学习快速入门教程（绝对通俗易懂！）【小土堆】