SEED数据集实战用PyTorch搭建EEG情绪识别模型从数据加载到训练脑电信号EEG情绪识别是脑机接口领域的前沿方向之一。SEED数据集作为公开可获取的高质量EEG数据集为研究者提供了标准化的实验数据。本文将手把手教你如何用PyTorch实现端到端的情绪识别模型开发。1. 环境准备与数据获取在开始之前确保已安装以下Python库pip install torch torchvision numpy pandas matplotlib scikit-learnSEED数据集可从上海交通大学BCMI实验室官网申请获取。下载后解压得到以下目录结构Preprocessed_EEG/ ├── 1/ │ ├── 1_20131130.mat │ └── ... ├── 2/ │ ├── 2_20131027.mat │ └── ... └── ...提示SEED数据集包含15名受试者在观看情感视频时的62通道EEG记录采样率200Hz已进行降噪和分段预处理。2. 数据加载与预处理2.1 自定义数据集类首先创建继承自torch.utils.data.Dataset的自定义类import torch from torch.utils.data import Dataset import scipy.io as sio import numpy as np class SEEDDataset(Dataset): def __init__(self, root_dir, subjectsNone, transformNone): self.root_dir root_dir self.transform transform self.data [] self.labels [] # 默认加载所有受试者数据 subjects subjects if subjects else range(1, 16) for subj in subjects: mat_files [f for f in os.listdir(f{root_dir}/{subj}) if f.endswith(.mat)] for mat_file in mat_files: mat_data sio.loadmat(f{root_dir}/{subj}/{mat_file}) eeg mat_data[data] # 形状为(62, timesteps, trials) label mat_data[labels].flatten() # 按试验分割数据 for trial in range(eeg.shape[2]): self.data.append(eeg[:, :, trial]) self.labels.append(label[trial]) def __len__(self): return len(self.data) def __getitem__(self, idx): sample self.data[idx] label self.labels[idx] if self.transform: sample self.transform(sample) return torch.FloatTensor(sample), torch.LongTensor([label])2.2 数据标准化与增强EEG数据通常需要以下预处理步骤带通滤波保留0.5-50Hz频段重参考转换为平均参考标准化通道级Z-score标准化from sklearn.preprocessing import StandardScaler class EEGTransform: def __call__(self, eeg): # 时间维度标准化 scaler StandardScaler() eeg scaler.fit_transform(eeg.T).T # 添加维度作为CNN输入通道 return eeg[np.newaxis, :, :]3. 模型架构设计3.1 混合CNN-LSTM网络结合CNN的局部特征提取和LSTM的时序建模能力import torch.nn as nn class EEGEmotionNet(nn.Module): def __init__(self, num_classes3): super().__init__() self.cnn nn.Sequential( nn.Conv2d(1, 16, kernel_size(1, 5), padding(0, 2)), nn.BatchNorm2d(16), nn.ELU(), nn.MaxPool2d((1, 2)), nn.Conv2d(16, 32, kernel_size(62, 1)), # 空间卷积 nn.BatchNorm2d(32), nn.ELU(), nn.Dropout(0.5) ) self.lstm nn.LSTM(input_size32, hidden_size64, num_layers2, bidirectionalTrue) self.classifier nn.Sequential( nn.Linear(128, 64), nn.ReLU(), nn.Dropout(0.5), nn.Linear(64, num_classes) ) def forward(self, x): # x形状: (batch, 1, 62, timesteps) x self.cnn(x) # (batch, 32, 1, timesteps//2) x x.squeeze(2).permute(2, 0, 1) # (timesteps//2, batch, 32) _, (h_n, _) self.lstm(x) h_n h_n.view(2, 2, -1, 64)[-1] # 取最后层双向隐藏状态 h_n h_n.permute(1, 0, 2).reshape(-1, 128) # (batch, 128) return self.classifier(h_n)3.2 模型参数分析层类型输出形状参数量说明Conv2d(16,62,200)96时间卷积MaxPool(16,62,100)-时间维度下采样Conv2d(32,1,100)31,744空间卷积LSTM(50,64)49,664双向LSTMLinear(64)8,256全连接层Linear(3)195输出层4. 训练流程与技巧4.1 交叉验证策略采用受试者独立的交叉验证from sklearn.model_selection import LeaveOneGroupOut def subject_cv_train(dataset, n_splits5): logo LeaveOneGroupOut() subjects [i//100 for i in range(len(dataset))] # 假设每个受试者100个样本 for fold, (train_idx, val_idx) in enumerate(logo.split( range(len(dataset)), groupssubjects)): train_set torch.utils.data.Subset(dataset, train_idx) val_set torch.utils.data.Subset(dataset, val_idx) train_loader DataLoader(train_set, batch_size32, shuffleTrue) val_loader DataLoader(val_set, batch_size32) model EEGEmotionNet().to(device) optimizer torch.optim.AdamW(model.parameters(), lr1e-3) criterion nn.CrossEntropyLoss() # 训练循环 for epoch in range(30): model.train() for x, y in train_loader: x, y x.to(device), y.to(device) optimizer.zero_grad() outputs model(x) loss criterion(outputs, y.squeeze()) loss.backward() optimizer.step() # 验证评估 model.eval() val_acc evaluate(model, val_loader) print(fFold {fold}, Epoch {epoch}, Val Acc: {val_acc:.2f})4.2 关键训练技巧学习率调度使用余弦退火scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max10, eta_min1e-5)标签平滑缓解过拟合criterion nn.CrossEntropyLoss(label_smoothing0.1)混合精度训练加速训练过程scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): outputs model(x) loss criterion(outputs, y.squeeze()) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()5. 模型评估与可视化5.1 评估指标除准确率外建议关注混淆矩阵分析各类别识别情况F1-score处理类别不平衡AUC-ROC评估模型区分能力from sklearn.metrics import classification_report def evaluate(model, loader): model.eval() all_preds, all_labels [], [] with torch.no_grad(): for x, y in loader: x x.to(device) outputs model(x) preds torch.argmax(outputs, dim1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(y.squeeze().cpu().numpy()) print(classification_report(all_labels, all_preds)) return accuracy_score(all_labels, all_preds)5.2 特征可视化使用t-SNE降维观察特征分布from sklearn.manifold import TSNE import matplotlib.pyplot as plt def visualize_features(model, loader): model.eval() features, labels [], [] with torch.no_grad(): for x, y in loader: x x.to(device) h model.cnn(x) h h.mean(dim[2,3]) # 全局平均池化 features.append(h.cpu()) labels.append(y.cpu()) features torch.cat(features).numpy() labels torch.cat(labels).numpy() tsne TSNE(n_components2) reduced tsne.fit_transform(features) plt.figure(figsize(10,8)) for i in range(3): # 假设3类情绪 plt.scatter(reduced[labelsi,0], reduced[labelsi,1], labelfClass {i}) plt.legend() plt.title(EEG Feature Visualization) plt.show()6. 实际应用建议在真实场景部署EEG情绪识别模型时有几个实用建议实时处理优化将模型转换为ONNX格式提升推理速度torch.onnx.export(model, dummy_input, eeg_model.onnx, input_names[eeg], output_names[emotion])领域自适应当应用于新受试者时可采用以下策略特征对齐使用CORAL等域适应方法微调在新数据上少量迭代训练多模态融合结合面部表情或生理信号提升鲁棒性class MultimodalModel(nn.Module): def __init__(self): super().__init__() self.eeg_branch EEGEmotionNet() self.face_branch ResNet18() self.fusion nn.Linear(256, 3) def forward(self, eeg, face): eeg_feat self.eeg_branch(eeg) face_feat self.face_branch(face) combined torch.cat([eeg_feat, face_feat], dim1) return self.fusion(combined)