简介本资源是一套基于脑电图EEG信号实现抑郁症智能辅助诊断的Python开源代码面向生物医学工程、人工智能医疗方向的研究者与高校研究生解决临床EEG数据建模难、图神经网络应用门槛高等实际问题。压缩包共4个文件含3个核心Python脚本分别实现Chebyshev图卷积模型构建、聚类指标计算及EEG数据预处理流程和1份README.md说明文档整体仅7KB轻量紧凑、即下即用。已有714人学习下载适合开展抑郁症脑机制研究、GCN在生理信号建模中的迁移实践或作为课程设计中图神经网络与医疗AI交叉课题的完整参考实现。代码结构清晰模块职责明确覆盖从原始EEG数据加载、图结构构建、SSPA-GCN模型训练到评估的全流程附带关键参数配置与可复现实验逻辑便于快速理解算法设计思想并进行二次开发。1. 为什么用 SSPA-GCN 做 EEG 抑郁症诊断不是玄学而是可复现的信号建模闭环你手头有一段 64 导联、250Hz 采样的静息态 EEG 数据想判断被试是否存在临床级抑郁倾向——但传统机器学习比如 SVM手工特征在跨被试泛化时 AUC 常掉到 0.65 以下而端到端 CNN 又对电极空间拓扑关系“视而不见”。这时候SSPA-GCNSpectral-Spatial Attention Graph Convolutional Network就不是论文里的炫技名词而是把脑电信号的频域特性、电极物理位置、通道间功能耦合三者强行捆在一起建模的务实方案。它不依赖 fMRI 或结构 MRI 辅助纯靠 EEG 原始电压序列 标准 10-20 系统电极坐标就能跑通模型结构里嵌了可学习的图注意力权重能自动识别前额叶-颞叶通路在 theta 波段的异常同步衰减——这正是临床文献反复验证的抑郁症生物标志物。适合神经工程方向的研究生快速验证假设也适合精神科数字医疗团队做二类医疗器械算法预研。注意它不是黑匣子诊断工具而是把医生已知的病理机制如 alpha 波不对称性、theta/gamma 耦合紊乱编码进图结构与注意力机制的可解释建模框架。2. 搭建 SSPA-GCN 训练环境从 EEG 数据预处理到图结构构建2.1 安装核心依赖与版本锁定策略SSPA-GCN 对 PyTorch 版本敏感尤其涉及torch_geometric的图卷积层。实测torch1.13.1cu117torch-geometric2.2.0组合在 NVIDIA A100 上训练最稳低于 1.12 会触发MessagePassing的 backward 内存泄漏高于 1.14 则torch-scatter编译失败。不要用pip install torch-geometric直接装——必须按官方文档指定源安装# 先卸载可能冲突的旧包 pip uninstall torch-scatter torch-sparse torch-cluster torch-spline-conv -y # 按 CUDA 版本选对应 wheel以 cu117 为例 pip install torch1.13.1cu117 torchvision0.14.1cu117 --extra-index-url https://download.pytorch.org/whl/cu117 pip install torch-geometric2.2.0 -f https://data.pyg.org/whl/torch-1.13.1cu117.html提示若用 CPU 环境把cu117替换为cpu但训练速度会降 5 倍以上建议至少用 RTX 3090 起步。scikit-learn1.2.2必须锁定新版1.3的StratifiedKFold会改变随机种子行为导致论文结果不可复现。2.2 EEG 数据标准化流程从 .edf 到 numpy 张量SSPA-GCN 输入要求是(N, C, T)形状张量N为样本数C为电极数默认 64T为时间点需统一截取为 2500 点即 10 秒 250Hz。关键步骤不是简单归一化而是保留病理信息的带通滤波 伪迹抑制import mne import numpy as np from scipy.signal import butter, filtfilt def preprocess_eeg(raw_path: str) - np.ndarray: # 1. 加载并重参考推荐平均参考非耳垂参考 raw mne.io.read_raw_edf(raw_path, preloadTrue) raw.set_eeg_reference(average) # 2. 0.5–45Hz 带通滤波去直流高频噪声保留 delta 到 gamma 全频段 b, a butter(4, [0.5, 45], btypebandpass, fsraw.info[sfreq]) raw._data filtfilt(b, a, raw._data, axis-1) # 3. ICA 去眼电/肌电仅对含明显伪迹的被试启用 ica mne.preprocessing.ICA(n_components20, random_state42) ica.fit(raw) raw ica.apply(raw, exclude[0, 1]) # 排除前两个成分通常为眼电 # 4. 截取 10 秒片段取中间段避开开头不稳定期 start_idx (raw.n_times - 2500) // 2 data raw.get_data()[:, start_idx:start_idx2500] # (C, T) return data.astype(np.float32) # 批量处理示例 eeg_list [] for f in [sub001.edf, sub002.edf]: eeg_list.append(preprocess_eeg(f)) X np.stack(eeg_list) # (N, C, T)逻辑说明mne是唯一能稳定解析多厂商 .edf 格式的库butter(4,...)用 4 阶巴特沃斯滤波器避免相位失真ICA 排除成分编号[0,1]是经验阈值——实际需用ica.plot_sources(raw)人工确认此处为自动化脚本简化。参数说明start_idx计算确保每段都是连续 10 秒避免拼接伪迹astype(np.float32)为后续 GPU 计算省显存。2.3 构建电极空间图用 10-20 系统坐标生成邻接矩阵SSPA-GCN 的核心是图结构——不是随便连边而是用欧氏距离加高斯核生成带权邻接矩阵A再通过k8近邻稀疏化import numpy as np from sklearn.metrics.pairwise import euclidean_distances # 10-20 系统 64 导联三维坐标单位cm已校准球面投影 electrode_coords np.array([ [-8.2, 5.1, 8.3], # Fp1 [ 8.2, 5.1, 8.3], # Fp2 [-5.4, 8.7, 5.2], # F7 [ 5.4, 8.7, 5.2], # F8 # ... 共 64 行完整坐标表见项目 data/electrode_64.npy ]) # 计算距离矩阵并高斯核加权 D euclidean_distances(electrode_coords) # (64, 64) sigma np.mean(D[D 0]) * 0.3 # 自适应带宽 A_dense np.exp(-D**2 / (2 * sigma**2)) # (64, 64) # k 近邻稀疏化每个节点只连最近的 8 个邻居 A_sparse np.zeros_like(A_dense) for i in range(A_dense.shape[0]): idx np.argsort(A_dense[i])[-9:] # 包含自身取 top-9 A_sparse[i, idx] A_dense[i, idx] A_sparse (A_sparse A_sparse.T) / 2 # 对称化 np.save(adj_matrix_64.npy, A_sparse) # 保存供模型加载逻辑说明sigma用距离均值的 30% 是经验值——太小导致图过于稀疏丢失长程连接太大则变成全连接图失去空间选择性k8来自论文实验在 64 导联下8 个邻居能覆盖额叶-顶叶-枕叶主要功能环路且图卷积层数控制在 3 层内不爆炸。参数说明坐标数据必须用electrode_64.npy非标准 10-10 系统否则电极映射错位会导致模型学不到真实脑网络。3. SSPA-GCN 模型实现三层图卷积 双路注意力机制3.1 模型主干Spectral-Spatial Attention BlockSSPA-GCN 不是简单堆叠 GCN而是将频域Spectral和空间Spatial注意力解耦设计。关键创新在SpectralAttention模块——它对每个电极的时间序列做 FFT 后在频点维度施加注意力权重而非传统 CNN 在通道维度加权import torch import torch.nn as nn import torch.nn.functional as F from torch_geometric.nn import GCNConv class SpectralAttention(nn.Module): def __init__(self, T: int, freq_bins: int 128): super().__init__() self.freq_bins freq_bins self.attention nn.Sequential( nn.Linear(T, 64), nn.ReLU(), nn.Linear(64, freq_bins) # 输出每个频点的权重 ) def forward(self, x: torch.Tensor) - torch.Tensor: # x: (B, C, T) - FFT to (B, C, freq_bins) X_fft torch.fft.rfft(x, nself.freq_bins, dim-1).abs() # 生成频域注意力权重 weights torch.sigmoid(self.attention(x.mean(dim1))) # (B, freq_bins) # 加权融合 out (X_fft * weights.unsqueeze(1)).sum(dim-1) # (B, C) return out class SSPA_GCN(nn.Module): def __init__(self, num_nodes64, in_features2500, hidden_dim128, num_classes2): super().__init__() self.spectral_att SpectralAttention(Tin_features) self.gcn1 GCNConv(num_nodes, hidden_dim) # 第一层节点特征聚合 self.gcn2 GCNConv(hidden_dim, hidden_dim) self.gcn3 GCNConv(hidden_dim, hidden_dim) self.spatial_att nn.MultiheadAttention(embed_dimhidden_dim, num_heads4, batch_firstTrue) self.classifier nn.Sequential( nn.Linear(hidden_dim, 64), nn.ReLU(), nn.Dropout(0.5), nn.Linear(64, num_classes) ) def forward(self, x: torch.Tensor, edge_index: torch.Tensor, edge_weight: torch.Tensor): # Step 1: Spectral attention → (B, C) spec_feat self.spectral_att(x) # (B, C) # Step 2: GCN layers → (B, C, hidden_dim) x_gcn spec_feat.unsqueeze(-1) # (B, C, 1) x_gcn F.relu(self.gcn1(x_gcn, edge_index, edge_weight)) x_gcn F.relu(self.gcn2(x_gcn, edge_index, edge_weight)) x_gcn self.gcn3(x_gcn, edge_index, edge_weight) # (B, C, hidden_dim) # Step 3: Spatial attention across electrodes x_att, _ self.spatial_att(x_gcn, x_gcn, x_gcn) # (B, C, hidden_dim) x_pool x_att.mean(dim1) # 全局平均池化 return self.classifier(x_pool)逻辑说明SpectralAttention的torch.fft.rfft只计算实数 FFT节省显存n128覆盖 0–64HzNyquist 频率的一半weights.unsqueeze(1)实现频点对每个电极的广播乘法GCN 层输入x_gcn是(B,C,1)因为 SSPA-GCN 将每个电极视为图节点时间维度已被频域压缩。参数说明hidden_dim128是平衡精度与显存的临界值——小于 64 时模型欠拟合AUC0.72大于 256 时梯度消失loss 不下降。3.2 数据加载器动态图构建与标签平衡SSPA-GCN 要求每个 batch 内edge_index和edge_weight与当前 batch 大小匹配。不能用静态图——必须在DataLoader中实时构建from torch_geometric.data import Data from torch.utils.data import Dataset, DataLoader class EEGGraphDataset(Dataset): def __init__(self, X: np.ndarray, y: np.ndarray, adj_matrix: np.ndarray): self.X torch.from_numpy(X).float() # (N, C, T) self.y torch.from_numpy(y).long() # (N,) self.adj torch.from_numpy(adj_matrix).float() # (C, C) def __len__(self): return len(self.X) def __getitem__(self, idx): x self.X[idx] # (C, T) y self.y[idx] # 动态构建图从邻接矩阵提取 edge_index 和 edge_weight adj_triu torch.triu(self.adj, diagonal1) # 上三角 edge_weight adj_triu[adj_triu 0].view(-1) edge_index torch.nonzero(adj_triu, as_tupleTrue) edge_index torch.stack(edge_index, dim0) # (2, E) # 转换为 PyG Data 对象 data Data(xx.T, edge_indexedge_index, edge_attredge_weight, yy) return data # 使用示例 dataset EEGGraphDataset(X_train, y_train, A_sparse) loader DataLoader(dataset, batch_size16, shuffleTrue, collate_fncollate_fn)逻辑说明collate_fn需自定义代码略因 PyG 默认collate_fn无法处理edge_attrx.T是关键——SSPA-GCN 要求节点特征是(T, C)但输入张量是(C, T)故转置torch.triu(..., diagonal1)避免自环电极不连自己符合脑网络建模惯例。参数说明batch_size16是显存安全上限——RTX 3090 下32会 OOM8则训练震荡。4. 训练与验证损失函数设计、早停策略与跨被试评估4.1 抑郁症诊断特有的损失函数Focal Loss 类别权重抑郁症数据集天然不平衡健康:患者 ≈ 3:1直接用CrossEntropyLoss会导致模型偏向多数类。SSPA-GCN 论文采用 Focal Loss 改进但需调整gamma参数class FocalLoss(nn.Module): def __init__(self, alpha1, gamma2, reductionmean): super().__init__() self.alpha alpha self.gamma gamma self.reduction reduction def forward(self, inputs, targets): ce_loss F.cross_entropy(inputs, targets, reductionnone) pt torch.exp(-ce_loss) focal_weight (1 - pt) ** self.gamma loss focal_weight * ce_loss if self.reduction mean: return loss.mean() elif self.reduction sum: return loss.sum() else: return loss # 初始化损失函数alpha 根据类别频率计算 from sklearn.utils.class_weight import compute_class_weight class_weights compute_class_weight(balanced, classesnp.unique(y_train), yy_train) criterion FocalLoss(alphatorch.tensor(class_weights, dtypetorch.float32).to(device))逻辑说明gamma2是论文默认值但实测在 EEG 数据上gamma1.5更稳——过高gamma会使难样本梯度爆炸compute_class_weight比手动设alpha[1,3]更鲁棒因不同数据集不平衡比差异大。参数说明reductionmean必须保持否则loss.backward()会报错。4.2 跨被试验证协议Leave-One-Subject-OutLOSO临床场景要求模型不接触目标被试数据SSPA-GCN 必须用 LOSO 而非随机 K-Foldfrom sklearn.model_selection import LeaveOneGroupOut from tqdm import tqdm def train_loso(model, X, y, groups, device): logo LeaveOneGroupOut() results {auc: [], acc: [], f1: []} for train_idx, test_idx in logo.split(X, y, groups): # 分割数据 X_train, X_test X[train_idx], X[test_idx] y_train, y_test y[train_idx], y[test_idx] # 训练模型含早停 model train_one_fold(model, X_train, y_train, device) # 测试 y_pred predict(model, X_test, device) results[auc].append(roc_auc_score(y_test, y_pred[:, 1])) results[acc].append(accuracy_score(y_test, y_pred.argmax(1))) results[f1].append(f1_score(y_test, y_pred.argmax(1))) return results # groups 参数示例[0,0,0,1,1,2,2,2,2,...] 表示每个样本所属被试 ID groups np.array([0]*30 [1]*28 [2]*32 ...) # 每个被试约 30 个样本 results train_loso(model, X, y, groups, device) print(fLOSO AUC: {np.mean(results[auc]):.3f} ± {np.std(results[auc]):.3f})逻辑说明LeaveOneGroupOut确保测试集完全独立于训练集无数据泄露groups必须是被试 ID 数组predict()函数需返回概率而非硬标签因roc_auc_score需要置信度。参数说明LOSO 循环次数 被试数若超 50 人建议用tqdm显示进度条避免误判卡死。4.3 早停与学习率调度基于验证集 AUC 的动态调整SSPA-GCN 训练易过拟合需严格早停。但不用val_loss——用val_auc因为抑郁症诊断更关注排序能力best_auc 0.0 patience 15 trigger_times 0 for epoch in range(100): train_loss train_epoch(model, train_loader, optimizer, criterion, device) val_auc validate_auc(model, val_loader, device) # 返回 AUC 值 if val_auc best_auc: best_auc val_auc torch.save(model.state_dict(), best_sspa_gcn.pth) trigger_times 0 else: trigger_times 1 # 学习率衰减当 AUC 连续 10 轮不升lr * 0.5 if trigger_times 10 and epoch % 10 0: for param_group in optimizer.param_groups: param_group[lr] * 0.5 if trigger_times patience: print(fEarly stopping at epoch {epoch}) break逻辑说明validate_auc函数内部用sklearn.metrics.roc_auc_score计算而非model.eval()后的losspatience15是经验值——小于 10 会早停AUC 波动正常大于 20 则过拟合验证 AUC 开始掉。参数说明学习率衰减触发条件是trigger_times 10而非10避免抖动干扰。5. 避坑指南SSPA-GCN 实战中 4 个血泪教训5.1 现象训练 loss 下降但验证 AUC 停滞在 0.55模型学不会任何判别模式原因电极坐标文件electrode_64.npy与你的 EEG 数据实际导联顺序不匹配。例如你的 .edf 文件按Fp1,Fp2,F3,F4,...排列但坐标文件按Fp1,F3,F5,...奇偶分组排列导致图结构完全错误。解决用mne.channels.make_standard_montage(standard_1005)生成标准坐标再用raw.get_montage().get_positions()[ch_pos]提取实际电极位置与electrode_64.npy逐点比对欧氏距离误差 1cm 即需重排。5.2 现象GPU 显存占用 100% 但 batch_size1 仍 OOM原因torch-geometric的GCNConv在反向传播时缓存中间变量当hidden_dim 128且num_nodes64时显存需求呈平方增长。解决在GCNConv后插入torch.cuda.empty_cache()或改用ChebConv需预计算切比雪夫多项式或降低hidden_dim至 96 并增加 GCN 层数如 4 层 × 96。5.3 现象LOSO 验证 AUC 方差极大±0.15结果不可信原因部分被试 EEG 数据质量差如运动伪迹未清除其测试集 AUC 拉低整体均值。但LeaveOneGroupOut无法剔除坏被试。解决预处理阶段加入信噪比SNR筛选——计算每个被试所有通道的var(signal)/var(noise)SNR 5dB 的被试整例剔除并在论文方法部分注明剔除比例。5.4 现象模型输出概率全部接近 [0.5, 0.5]分类无区分度原因SpectralAttention的torch.fft.rfft输入张量未归一化FFT 幅度值过大导致 sigmoid 饱和。解决在SpectralAttention.forward中对输入x做x F.normalize(x, p2, dim-1)L2 归一化或在 FFT 前x x / x.std(dim-1, keepdimTrue)。6. 模型可解释性落地用 Grad-CAM 定位抑郁症关键脑区与频段SSPA-GCN 的价值不仅在于准确率更在于它能回答“模型为什么认为这是抑郁症”——这需要把图注意力权重和频域注意力可视化。我一般不用论文里复杂的 GNNExplainer而是用轻量级 Grad-CAM 变体直接作用于SpectralAttention模块def spectral_cam(model, x: torch.Tensor, target_class: int 1): model.eval() x.requires_grad_(True) # 前向传播获取频域注意力权重 spec_feat model.spectral_att(x) # (B, C) logits model.classifier(spec_feat) # (B, 2) # 计算目标类别的梯度 loss logits[:, target_class].sum() loss.backward() # 获取梯度对输入 x 的梯度 gradients x.grad # (B, C, T) # 全局平均池化梯度得到权重 weights gradients.mean(dim[0, 2]) # (C,) —— 每个电极的重要性 # 可视化热力图叠加在 10-20 系统脑图上 plot_topomap(weights.numpy(), electrode_coords, ch_typeeeg, showTrue, titlefSpectral CAM for class {target_class}) # 调用示例对单个样本 x_sample X_test[0:1] # (1, C, T) spectral_cam(model, torch.from_numpy(x_sample).to(device))逻辑说明gradients.mean(dim[0,2])是关键——dim0对 batch 平均虽为 1 但保持维度dim2对时间点平均得到每个电极的综合重要性plot_topomap用mne.viz.plot_topomap实现自动映射到头皮三维坐标。参数说明target_class1对应抑郁症类别若模型输出logits未 softmax需先F.softmax(logits, dim1)。更进一步我把 Grad-CAM 结果和临床知识对齐取权重最高的 5 个电极如 Fp1, F3, F7, T3, C3再分析其SpectralAttention输出的频点权重——发现 theta 波段4–8Hz权重占比超 65%这与文献报道的前额叶 theta 功率升高一致。这种闭环验证让我敢把模型交给医生试用而不是当黑盒交差。最后说个血泪习惯每次跑完 LOSO我必用shap库对spec_feat做全局解释生成summary_plot。如果 SHAP 值分布和 Grad-CAM 热力图趋势相反说明模型学到的是数据集偏差如某设备采集的伪迹立刻停用该 fold。技术可以炫但临床决策容不得半点侥幸。希望帮到你。本文还有配套的精品资源点击获取