1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
| """
MVMD-FCASNN 核心实现
论文: Yao, Gu, Wang (2026) Medical & Biological Engineering & Computing
"""
import torch
import torch.nn as nn
import numpy as np
from scipy.signal import hilbert
def mvmd_decomposition(eeg_signal, num_imfs=4, max_iter=100):
"""
多元变分模态分解 (简化版)
与传统VMD逐通道分解不同，MVMD联合分解所有通道，
保持频率对齐和空间关系
Args:
eeg_signal: (C, T) 多通道EEG信号
num_imfs: IMF数量
max_iter: 最大迭代次数
Returns:
imfs: (num_imfs, C, T) 各阶IMF
"""
C, T = eeg_signal.shape
imfs = np.zeros((num_imfs, C, T))
for k in range(num_imfs):
center_freq = (k + 1) / (num_imfs + 1) * 0.5
from scipy.signal import butter, filtfilt
bw = 0.05
low = max(center_freq - bw, 0.001)
high = min(center_freq + bw, 0.499)
b, a = butter(4, [low, high], btype='band')
for ch in range(C):
imfs[k, ch] = filtfilt(b, a, eeg_signal[ch])
return imfs
class FrequencyChannelAttentionSNN(nn.Module):
"""
FCASNN: 频率+通道注意力脉冲神经网络
1. IMF注意力: 选择重要频段 (theta/alpha/beta)
2. 通道注意力: 选择重要脑区 (额叶/中央)
3. SNN: 类脑脉冲网络，事件驱动+时序保留
"""
def __init__(self,
num_channels: int = 32,
num_imfs: int = 4,
time_steps: int = 1000,
hidden_dim: int = 128):
super().__init__()
self.imf_attention = nn.Sequential(
nn.Linear(num_channels * time_steps, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, num_imfs),
nn.Softmax(dim=-1)
)
self.channel_attention = nn.Sequential(
nn.Linear(num_imfs * time_steps, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, num_channels),
nn.Sigmoid()
)
self.sn1 = LIFNeuron(num_imfs * num_channels, hidden_dim)
self.sn2 = LIFNeuron(hidden_dim, hidden_dim)
self.sn3 = LIFNeuron(hidden_dim, 2)
self.fc_out = nn.Linear(2, 2)
def forward(self, imfs):
"""
Args:
imfs: (batch, num_imfs, channels, time)
Returns:
logits: (batch, 2) 疲劳分类
"""
B, K, C, T = imfs.shape
imf_flat = imfs.view(B, K, C * T)
imf_weights = self.imf_attention(imf_flat)
weighted_imfs = imfs * imf_weights.view(B, K, 1, 1)
ch_flat = weighted_imfs.view(B, K * T, C)
ch_flat_perm = ch_flat.permute(0, 2, 1)
ch_weights = self.channel_attention(ch_flat_perm)
ch_weights = torch.sigmoid(torch.mean(weighted_imfs, dim=(1, 3)))
weighted = weighted_imfs * ch_weights.view(B, 1, C, 1)
x = weighted.view(B, -1)
x = self.sn1(x)
x = self.sn2(x)
x = self.sn3(x)
logits = self.fc_out(x)
return logits
class LIFNeuron(nn.Module):
"""
Leaky Integrate-and-Fire 脉冲神经元
模拟生物神经元:
- 积分: 累积输入
-泄漏: 时间衰减
- 发放: 达到阈值则发放脉冲
"""
def __init__(self, in_features, out_features,
threshold=1.0, leak=0.9):
super().__init__()
self.fc = nn.Linear(in_features, out_features)
self.threshold = threshold
self.leak = leak
def forward(self, x):
mem = self.fc(x)
mem = self.leak * mem
spikes = (mem >= self.threshold).float()
mem = mem * (1 - spikes)
return spikes
if __name__ == "__main__":
B, C, T = 8, 32, 1000
eeg = torch.randn(B, C, T)
eeg_np = eeg[0].numpy()
imfs = mvmd_decomposition(eeg_np, num_imfs=4)
print(f"EEG形状: {eeg.shape}")
print(f"IMF形状: {imfs.shape}")
model = FrequencyChannelAttentionSNN(num_channels=32, num_imfs=4)
imfs_tensor = torch.randn(B, 4, 32, T)
logits = model(imfs_tensor)
print(f"输出logits: {logits.shape}")
print("\n=== 论文实验结果 ===")
print(f"准确率: 98.81%")
print(f"-6dB SNR下: 仍保持高精度")
print(f"对比方法: EMD+SNN, VMD+SNN等均显著退化")
print("\n=== 关键频段 ===")
print("Theta (4-8Hz): 疲劳时增强")
print("Alpha (8-13Hz): 疲劳时增强 (额叶)")
print("Beta (13-30Hz): 疲劳时减弱 (注意力)")
|
