TorchEEG —— 基于 PyTorch 的模块化脑电分析

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

引言

2021 年，中国脑计划经过六年多的筹划终于尘埃落定，国家经费拨款超过 31 亿元人民币，涉及 59 个研究领域和方向的“脑科学与类脑研究”项目轰轰烈烈展开。在过去的几年中，笔者有幸见证了脑科学的飞速发展，心理学、神经科学、计算机科学领域的专家、学者和同学都纷至沓来，分享会和座谈会越来越多，给领域带来了新的生命力。

笔者所从事的研究属于脑科学中的一个很有趣的小分支，脑电分析。通过附着在头皮上的电极（金属盘）采集脑细胞活动时产生的生物电信号，以分析确定大脑的活动和变化。相比于其他侵入式的设备（需要通过手术将电极或传感器植入大脑），脑电的采集安全便捷，且相对廉价，在实用性上具有很大优势。

图1: 脑电采集和分析示意图，来自于 https://www.gastroepato.it/eeg.htm

正因如此，这几年脑电分析研究飞速发展，很多深度学习的顶会每年都会有 2-3 篇脑电分析的文章发表。尽管相关研究越来越多，由于庞杂多样的脑电分析设备，繁琐耗时的数据预处理等，目前社区仍没有形成的适用于深度学习的脑电分析工具集，很多优秀的工作也来不及整理并开源源代码。

与自然语言处理和计算机视觉领域“调包”的易用性相比，脑电分析领域的工具包显然是不够“保姆式”的。为了方便脑电分析领域的同学们，能够方便地使用 benchmark 数据集，复现相关工作，并专注于模型设计和修改，我们开源了 TorchEEG ，帮助大家一键完成深度学习中的“脏活累活” 。

注：TorchEEG 目前处于 Beta 测试阶段，小范围试用，如有疏漏敬请谅解。预计 7-8 月份完成测试阶段并发布正式版。

脑电分析的应用

目前，脑电分析已经被广泛应用于情感识别、运动想象、睡眠分期检测、癫痫检测、视觉分类等任务中。通过采集脑电信号，深度学习方法已经在帮助人们识别被试者的情感，正在想象的运动，睡眠的阶段，是否患有神经疾病，所见的内容等方面做出了突出贡献。

尽管如此，脑电采集的过程相对耗时，需要组织被试者长时间接受情感等刺激。因此，目前的开放数据集大多规模有限（被试数量大多不超过 100 人）。此外，不同设备、实验设置也导致数据集之间的差异较大，使用起来有一定障碍，难以协同。

为了帮助同学们在实验初期无痛上手一个脑电分析的研究领域，TorchEEG 参考相关文献中的实验设置对开放数据集进行了统一的处理和封装：如何读取脑电信号、如何对齐相关文献并进行基本处理、如何分离基线信号、如何进行样本划分、如何匹配描述信息等。

例如，在 TorchEEG 1.0.x 版本中，我们的开源计划主要集中在情感识别数据集上：

AMIGOS 数据集来自论文: AMIGOS: A dataset for affect, personality and mood research on individuals and groups .
DREAMER 数据集来自论文: DREAMER: A database for emotion recognition through EEG and ECG signals from wireless low-cost off-the-shelf devices .
SEED 数据集来自论文: Investigating critical frequency bands and channels for EEG-based emotion recognition with deep neural networks .
DEAP 数据集来自论文: DEAP: A database for emotion analysis; using physiological signals .
MAHNOB 数据集来自论文: A multimodal database for affect recognition and implicit tagging .

出于隐私保护和版权的考虑，TorchEEG 并不提供下载代理。请同学们根据文档中的介绍前往官方地址签署协议，并按照协议进行下载和引用。我们不鼓励使用任何形式非官方下载渠道。在下载完成后，填写下载文件路径即可一键使用数据集，以下是一个例子：

from torcheeg.datasets import DEAPDataset
dataset = DEAPDataset(
    io_path=f'./deap',
    root_path='./data_preprocessed_python',
    online_transform=transforms.Compose(
        [transforms.To2d(),
         transforms.ToTensor()]),
    label_transform=transforms.Compose([
        transforms.Select(['valence', 'arousal']),
        transforms.Binary(5.0),
        transforms.BinariesToCategory()
    ]))
print(dataset[0])
# 数据集中的第一个样本
# 包含一个脑电信号，被表征为 torch.Tensor (1,32,128) 的张量,
# 包含一个标签，被表征为 int 的整数，0 表示情绪的效价低（消极），1 表示情绪的效价高（积极）

相关源码部分添加了详尽的注释，对于希望深入定制的同学，可以仔细阅读源码并进行修改（尽管笔者认为提供的设置选项已经足以覆盖大部分相关文献的实验设置）。同时，TorchEEG 欢迎任何人提交 issue、PR～

脑电分析的实验设置

与计算机视觉和自然语言处理领域的研究不同，脑电分析任务受到数据稀少、域鸿沟严重等因素的影响，实验设置对最终的性能影响非常大。其中，具有不同研究重点的研究工作对数据集划分的设计千差万别。TorchEEG 收集了相关文献中的常用设置，以便同学们能够根据自己的研究目标，合理选择并与相关 benchmark 进行对比：

传统的数据划分：KFoldDataset , train_test_split_dataset ,KFoldTrial , train_test_split_trial

传统的数据KFoldDataset, train_test_split_dataset 划分直接遵循了 KFold、比例划分等传统的交叉验证标准，在不进行任何分组的情况下对整个数据集进行随机抽样。笔者认为在这种划分方式下，训练样本和测试样本数量均较为充足，适合研究初期度量模型的区分能力使用。

但是，这些划分不进行任何分组，导致相邻脑电信号在很大概率上被分别划分在训练集合和测试集合中。相邻脑电信号的分布一致性保证了训练集合和测试结合分布的一致性。在真实应用场景中，我们往往无法找到这样的训练样本。因此，一些工作考虑了KFoldTrial, train_test_split_trial 交叉验证，抽取每次实验的连续若干时间的样本作为训练样本，其他样本作为测试样本，从而在一定程度上测试模型在一个实验不同时段上的泛化性。

被试依赖的数据划分：KFoldTrialPerSubject , train_test_split_trial_per_subject

# 一个简单的使用案例
from torcheeg.model_selection import KFoldTrial
cv = KFoldTrial(n_splits=5, shuffle=False, split_path='./split')
dataset = DEAPDataset(
    io_path=f'./deap',
    root_path='./data_preprocessed_python',
    online_transform=transforms.Compose(
        [transforms.To2d(),
         transforms.ToTensor()]),
    label_transform=transforms.Compose([
        transforms.Select(['valence', 'arousal']),
        transforms.Binary(5.0),
        transforms.BinariesToCategory()
    ]))
for train_dataset, test_dataset in cv.split(dataset):
    train_loader = DataLoader(train_dataset)
    test_loader = DataLoader(test_dataset)
    ...

跨被试的数据划分 :LeaveOneSubjectOut

目前较为流行的研究手段。将一个被试的所有样本作为测试样本，其他被试的所有样本作为训练样本。从而模拟在被试未知时，脑电分析的性能。能够在一定程度上测试模型在不同被试上的泛化性。

脑电分析中的难点

在实际应用中，脑电中所隐含的大脑活动特征并不是十分明显。需要我们客服克服一些困难才能够实现良好的脑电分析性能：

目前，笔者认为：

问题 1 的解决主要集中在特征工程上，通过去基线，去眼电等方式人工地提高脑电信号的信噪比，或使用预先定义对任务目标具有区分力的差分熵等特征进行特征提取。对于这一部分内容， TorchEEG 提供了丰富的变换器，使用者可以像在 CV 领域中调用函数对图像进行变换一样简单的对脑电信号实施特征工程。为了实现更高效的训练，我们提供了离线的多进程变换接口，通过参数设置即可使用多个 CPU 离线完成特征提取并预存供训练使用。

适用于特征工程 : BandDifferentialEntropy , BandPowerSpectralDensity , BandMeanAbsoluteDeviation ,BandKurtosis ,BandSkewness ,Concatenate
适用于常规预处理 :PickElectrode ,MeanStdNormalize ,MinMaxNormalize ,BaselineRemoval
适用于卷积神经网络（CNN）的预处理:To2d,ToGrid ,ToInterpolatedGrid
适用于图卷积神经网络（GNN）的预处理:ToG
适用于数据扩增 :Resize ,RandomNoise ,RandomMask
适用于标签变换 :Select ,Binary ,BinariesToCategory

from torcheeg import transforms
from torcheeg.datasets import AMIGOSDataset
from torcheeg.datasets.constants.emotion_recognition.deap import \
    AMIGOS_CHANNEL_LOCATION_DICT
dataset = AMIGOSDataset(
    io_path=f'./amigos',
    root_path='./data_preprocessed',
    # offline_transform, 预处理方法:
    # 离线预处理只只会运行一次
    # 运行结果缓存在本地，再次
    # 使用直接读取
    offline_transform=transforms.Compose([
        transforms.BandDifferentialEntropy(),
        transforms.ToGrid(AMIGOS_CHANNEL_LOCATION_DICT)
    ]),
    # online_transform, 在线预处理方法:
    # 每次索引都会运行
    online_transform=transforms.ToTensor(),
    label_transform=transforms.Compose([
        transforms.Select('valence'),
        transforms.Binary(5.0),
    ]),
    # 离线预处理支持高性能的多进程并行
    num_worker=4)

问题 2 的解决主要集中在模型设计上，通过不同模块显式的对时空信息进行建模，综合完成脑电信号的分析。对于这一部分内容， TorchEEG 提供了丰富的模型实现，无论是基于 CNN、GNN 还是流行的 Transformer 架构，都可以通过简单的调用和修改开始新的研究。

基于卷积神经网络（CNN）的方法：EEGNet,FBCCNN ,MTCNN ,STNet ,TSCeption ,CCNN
基于图卷积神经网络（GNN）的方法：DGCNN,RGNN
基于注意力机制（Transformer）的方法：SimpleViT

from torcheeg.models import EEGNet, FBCCNN
model = EEGNet().cuda()
model = FBCCNN().cuda()
...

问题 3 和问题 4 的解决主要集中在学习策略、损失函数的设计上，当前不同迁移学习、域泛化的思路有效地提升了脑电分析的性能。对于这一部分内容， TorchEEG 提供了丰富的样例，在筛选出鲁棒的相关工作后尽可能地还原并调优相关工作中提出的策略（目前仍在认真对齐性能中，敬请期待）。

脑电分析中的端到端建模

下面我们将以卷积神经网络和图卷积网络为例，分别展示两个端到端的建模过程。

使用卷积神经网络的案例（完整代码请见 github ）

首先，使用 pip 安装 TorchEEG 包：

pip install torcheeg
# # 如果想要安装未发布的最新版，也可以直接考虑从 github 安装
# pip install git+https://github.com/tczhangzhi/torcheeg.git

接着，使用 BandDifferentialEntropy 对 EEG 的每个电极的不同频段（通常为。alpha、beta、gamma、theta）提取差分熵特征，并使用 ToGrid 将电极映射到的网格上。

from torcheeg import transforms
from torcheeg.datasets import DEAPDataset
from torcheeg.datasets.constants.emotion_recognition.deap import \
    DEAP_CHANNEL_LOCATION_DICT
dataset = DEAPDataset(
    io_path=f'./tmp_out/deap',
    root_path='./tmp_in/data_preprocessed_python',
    offline_transform=transforms.Compose([
        transforms.BandDifferentialEntropy(
            apply_to_baseline=True),
        transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT,
                          apply_to_baseline=True)
    ]),
    online_transform=transforms.Compose(
        [transforms.BaselineRemoval(),
         transforms.ToTensor()]),
    label_transform=transforms.Compose([
        transforms.Select('valence'),
        transforms.Binary(5.0),
    ]),
    num_worker=4)

其中，子频段的差分熵描述了脑电信号中不同频率分量的时序变化的规律性，能够有效地将高维度、低信噪比的时序信号（一维时序向量）转变为若干个子频段的描述值（标量值），对于情感等人脑活动具有较强区分力。

而将电极映射到的网格上可以有效地为 CNN 模型提供电极间相对的位置信息。这里，电极在人脑上的相对位置被近似映射到一个网格上，使得相邻电极保持相邻。通常 10-20 脑电采集系统中标定的电极位置关系被用于网格映射。TorchEEG 提供了丰富的可视化接口（基于MNE 和matploitlab），最终的预处理效果如下，该 Tensor 将被作为神经网络的输入进行分析：

from torcheeg import transforms
from torcheeg.datasets import DEAPDataset
from torcheeg.datasets.constants.emotion_recognition.deap import \
    DEAP_CHANNEL_LOCATION_DICT
from torcheeg.utils import plot_3d_tensor
dataset = DEAPDataset(
    io_path=f'./tmp_out/deap',
    root_path='./tmp_in/data_preprocessed_python',
    offline_transform=transforms.Compose([
        transforms.BandDifferentialEntropy(
            apply_to_baseline=True),
        transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT,
                          apply_to_baseline=True)
    ]),
    online_transform=transforms.Compose(
        [transforms.BaselineRemoval(),
         transforms.ToTensor()]),
    label_transform=transforms.Compose([
        transforms.Select('valence'),
        transforms.Binary(5.0),
    ]),
    num_worker=4)
eeg = dataset[0][0]
plot_3d_tensor(eeg)

图2: 预处理的最终结果，该表征将作为 CNN 的输入使用（为了保护被试者隐私，这里展示了随机生成的脑电信号）

当然，TorchEEG 也支持以脑电地形图的形式进行展示：

from torcheeg import transforms
from torcheeg.datasets import DEAPDataset
from torcheeg.datasets.constants.emotion_recognition.deap import \
    DEAP_CHANNEL_LIST
from torcheeg.utils import plot_feature_topomap
dataset = DEAPDataset(
    io_path=f'./tmp_out/deap',
    root_path='./tmp_in/data_preprocessed_python',
    offline_transform=transforms.BandDifferentialEntropy(
        apply_to_baseline=True),
    online_transform=transforms.Compose(
        [transforms.BaselineRemoval(),
         transforms.ToTensor()]),
    label_transform=transforms.Compose([
        transforms.Select('valence'),
        transforms.Binary(5.0),
    ]),
    num_worker=4)
eeg = dataset[0][0]
plot_feature_topomap(
    eeg,
    channel_list=DEAP_CHANNEL_LIST,
    feature_list=["theta", "alpha", "beta", "gamma"])

图3: 不同子频段差分熵的脑电地形图（为了保护被试者隐私，这里展示了随机生成的脑电信号）

随后，使用 KFoldDataset 对数据集进行 kfold 划分，得到训练集合和测试集合。

from torcheeg.model_selection import KFoldDataset
k_fold = KFoldDataset(n_splits=10,
                      split_path=f'./tmp_out/split',
                      shuffle=True,
                      random_state=42)

接着，我们参照CCNN 的设计，建立一个小型但有效的 CNN 网络结构。

import torch
import torch.nn as nn
class CNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Sequential(
            nn.ZeroPad2d((1, 2, 1, 2)),
            nn.Conv2d(4, 64, kernel_size=4, stride=1),
            nn.ReLU())
        self.conv2 = nn.Sequential(
            nn.ZeroPad2d((1, 2, 1, 2)),
            nn.Conv2d(64, 128, kernel_size=4, stride=1),
            nn.ReLU())
        self.conv3 = nn.Sequential(
            nn.ZeroPad2d((1, 2, 1, 2)),
            nn.Conv2d(128, 256, kernel_size=4, stride=1),
            nn.ReLU())
        self.conv4 = nn.Sequential(
            nn.ZeroPad2d((1, 2, 1, 2)),
            nn.Conv2d(256, 64, kernel_size=4, stride=1),
            nn.ReLU())
        self.lin1 = nn.Linear(9 * 9 * 64, 1024)
        self.lin2 = nn.Linear(1024, 2)
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = x.flatten(start_dim=1)
        x = self.lin1(x)
        x = self.lin2(x)
        return

网络的主体部分由 4 个卷积层构成，每个卷积层通过对邻域的卷积分析局部相邻通道的差异并提取和大脑活动类别相关的模式特征：

图4: CCNN 中介绍的网络结构，这里使用的网络结构相似，但训练参数略有不同（最终的性能报告请参考原文，这里的案例建议作为实验起始点使用）

最后，使用交叉熵作为损失函数对模型进行训练。

from torch.utils.data.dataloader import DataLoader
for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
    model = CNN().to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
    train_loader = DataLoader(train_dataset,
                              batch_size=batch_size,
                              shuffle=True)
    val_loader = DataLoader(val_dataset,
                            batch_size=batch_size,
                            shuffle=False)
    epochs = 50
    for t in range(epochs):
        print(
            f"Epoch {t+1}
-------------------------------")
        train(train_loader, model, loss_fn, optimizer)
        valid(val_loader, model, loss_fn)
    print("Done!")

使用图卷积神经网络的案例（完整代码请见 github ）

首先，使用 pip 安装 TorchEEG 包：

pip install torcheeg
# # 如果想要安装未发布的最新版，也可以直接考虑从 github 安装
# pip install git+https://github.com/tczhangzhi/torcheeg.git

接着，使用 ToG 对 EEG 样本建立图表征。这里，电极通常对应于图结构中的节点，电极之间的关联关系被定义为边和边上的权值。通常考虑的关联关系有空间相邻和功能链接等。这里以空间相邻关系为例建立图结构，即邻接矩阵中的每个值表示两个对应电极是否在 10-20 系统中相邻，相邻为 1，不相邻为 0.

值的注意的是，笔者倾向于认为图卷积网络在节点角度的数据上表现较好（即电极数量较多，至少为 62 导），因此这里展示了一个 SEED 数据集上的使用案例。

from torcheeg import transforms
from torcheeg.datasets import SEEDDataset
from torcheeg.datasets.constants.emotion_recognition.seed import \
    SEED_ADJACENCY_MATRIX
dataset = SEEDDataset(
    io_path=f'./tmp_out/seed',
    root_path='./tmp_in/Preprocessed_EEG',
    offline_transform=transforms.BandDifferentialEntropy(),
    online_transform=transforms.ToG(SEED_ADJACENCY_MATRIX),
    label_transform=transforms.Compose([
        transforms.Select('emotion'),
        transforms.Lambda(lambda x: x + 1),
    ]),
    num_worker=8)

随后，练习使用 KFoldTrialPerSubject 数据划分方式，得到训练集合和测试集合。

from torcheeg.model_selection import KFoldTrialPerSubject
k_fold = KFoldTrialPerSubject(n_splits=10,
                              split_path=f'./tmp_out/split',
                              shuffle=False)

为了便于同学们利用现有开源工作开展研究，本文在建图后将脑电样本转化为torch-geometric 支持的 Data 类型。以便同学们基于模块化的torch-geometric 轻松建立模型。笔者基于GAT 提供了一个简单的案例，在开展研究的初期，同学们不妨尝试使用torch-geometric 中其他的卷积层进行实验，并分析不同卷积在 EEG 分析上的优劣以得到 insight 开展深入的自主设计。

from torch_geometric.nn import GATConv, global_mean_pool
class GNN(torch.nn.Module):
    def __init__(self,
                 in_channels=4,
                 num_layers=3,
                 hid_channels=64,
                 num_classes=3):
        super().__init__()
        self.conv1 = GATConv(in_channels, hid_channels)
        self.convs = torch.nn.ModuleList()
        for _ in range(num_layers - 1):
            self.convs.append(GATConv(hid_channels,
                                      hid_channels))
        self.lin1 = Linear(hid_channels, hid_channels)
        self.lin2 = Linear(hid_channels, num_classes)
    def reset_parameters(self):
        self.conv1.reset_parameters()
        for conv in self.convs:
            conv.reset_parameters()
        self.lin1.reset_parameters()
        self.lin2.reset_parameters()
    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch
        x = F.relu(self.conv1(x, edge_index))
        for conv in self.convs:
            x = F.relu(conv(x, edge_index))
        x = global_mean_pool(x, batch)
        x = F.relu(self.lin1(x))
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.lin2(x)
        return x

最后，使用交叉熵作为损失函数对模型进行训练。

from torch.utils.data.dataloader import DataLoader
for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
    model = CNN().to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
    train_loader = DataLoader(train_dataset,
                              batch_size=batch_size,
                              shuffle=True)
    val_loader = DataLoader(val_dataset,
                            batch_size=batch_size,
                            shuffle=False)
    epochs = 50
    for t in range(epochs):
        print(f"Epoch {t+1}
-------------------------------")
        train(train_loader, model, loss_fn, optimizer)
        valid(val_loader, model, loss_fn)
    print("Done!")

后记

由于篇幅限制，很多优秀的数据集和相关工作无法进行介绍。例如，运动想象和睡眠分期在近几年受到了广泛的关注，华人学者们已经在社区形成了相当的影响力。基于 2D 平面的卷积神经网络很早就取得了突出进展，CNN 与 LSTM 或与 GNN 的 hybrid 架构日趋受到欢迎。近年来，Transformer 也在 EEG 分析上取得了瞩目的成绩等等。TorchEEG 正在向这些为领域做出卓越贡献的相关工作和学者学习，并力图对齐、调优官方实现。TorchEEG 也欢迎同学们指出不足提出期望，多多 commit issue 和 PR～

最后，非常感谢我的两位指导老师，刘老师和钟老师的鼓励、安排、指导和提供的大量资源，尽管笔者磨磨蹭蹭马马虎虎，两位老师仍不抛弃不放弃，故余虽愚，卒获有所闻。也很感谢百忙之中答疑的同行，尽管笔者能力有限，但愿为社区发展尽绵薄之力！

引言

脑电分析的应用

脑电分析的实验设置

脑电分析中的难点

脑电分析中的端到端建模

使用卷积神经网络的案例（完整代码请见 github ）

使用图卷积神经网络的案例（完整代码请见 github ）

后记

Be First to Comment

发表回复取消回复

TorchEEG —— 基于 PyTorch 的模块化脑电分析

引言

脑电分析的应用

脑电分析的实验设置

脑电分析中的难点

脑电分析中的端到端建模

使用卷积神经网络的案例（完整代码请见 github ）

使用图卷积神经网络的案例（完整代码请见 github ）

后记

Be First to Comment

发表回复 取消回复

发表回复取消回复