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
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
| """
MVMD-FCASNN: 大脑启发疲劳检测框架
基于 Medical & Biological Engineering & Computing, 2026 复现
核心架构:
1. MVMD: 多变量变分模态分解
- 联合分解所有EEG通道
- 保持跨通道空间关系
- 对噪声鲁棒（带宽约束优化）
2. FCASNN: 频率+通道注意力脉冲神经网络
- IMF注意力: 自适应选择信息量最大的频率分量
- 通道注意力: 加权重要脑区（前额叶、中央区）
- SNN: 事件驱动，低功耗，保留时序结构
座舱应用迁移:
- 起重机操作员 → 驾驶员
- 建筑工地噪声 → 车内噪声（发动机/路面/空调）
- EEG头环 → 座椅内置EEG电极（未来）
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Tuple, List
class MVMD(nn.Module):
"""
Multivariate Variational Mode Decomposition (可学习版)
原始MVMD是优化算法，这里用可学习版本：
- 替代固定优化为可训练卷积
- 保留联合多通道分解的核心思想
"""
def __init__(self, num_channels: int = 14, num_imfs: int = 5):
super().__init__()
self.num_channels = num_channels
self.num_imfs = num_imfs
self.modal_filters = nn.Parameter(
torch.randn(num_imfs, num_channels, 64)
)
self.bandwidth = nn.Parameter(
torch.ones(num_imfs) * 0.1
)
def forward(self, eeg: torch.Tensor) -> torch.Tensor:
"""
Args:
eeg: (B, C, T) 多通道EEG信号
Returns:
imfs: (B, num_imfs, C, T) 分解后的模态分量
"""
B, C, T = eeg.shape
imfs = []
for i in range(self.num_imfs):
filt = self.modal_filters[i].unsqueeze(1)
imf = F.conv1d(
eeg, filt,
padding=32,
groups=C
)
imfs.append(imf)
return torch.stack(imfs, dim=1)
class IMFAttention(nn.Module):
"""
IMF频率注意力机制
自适应评估每个IMF分量的信息量
自动加权疲劳相关的频段（theta, alpha, beta）
"""
def __init__(self, num_imfs: int = 5, hidden_dim: int = 64):
super().__init__()
self.scorer = nn.Sequential(
nn.Linear(num_imfs, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, num_imfs),
nn.Sigmoid()
)
def forward(self, imfs: torch.Tensor) -> torch.Tensor:
"""
Args:
imfs: (B, num_imfs, C, T)
Returns:
weighted_imfs: (B, num_imfs, C, T)
"""
energy = imfs.pow(2).mean(dim=(2, 3))
weights = self.scorer(energy)
weighted = imfs * weights.unsqueeze(-1).unsqueeze(-1)
return weighted
class ChannelAttention(nn.Module):
"""
通道注意力机制（Squeeze-Excitation风格）
自适应加权脑区：
- 前额叶（F3, F4, Fz）: 认知控制 → 高权重
- 中央区（C3, C4, Cz）: 运动感知 → 中权重
- 枕叶（O1, O2）: 视觉处理 → 低权重（疲劳时）
"""
def __init__(self, num_channels: int = 14, reduction: int = 4):
super().__init__()
self.squeeze = nn.AdaptiveAvgPool1d(1)
self.excite = nn.Sequential(
nn.Linear(num_channels, num_channels // reduction),
nn.ReLU(),
nn.Linear(num_channels // reduction, num_channels),
nn.Sigmoid()
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x: (B, C, T)
Returns:
weighted: (B, C, T)
"""
squeezed = self.squeeze(x).squeeze(-1)
weights = self.excite(squeezed)
return x * weights.unsqueeze(-1)
class SpikingNeuronLayer(nn.Module):
"""
脉冲神经元层 (LIF模型)
Leaky Integrate-and-Fire 脉冲神经元:
- 事件驱动: 仅在输入变化时响应
- 低功耗: 适合可穿戴设备
- 保留时序结构: 毫秒级精度
"""
def __init__(self, in_features: int, out_features: int,
tau: float = 20.0, threshold: float = 1.0):
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.tau = tau
self.threshold = threshold
self.weight = nn.Linear(in_features, out_features, bias=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
LIF脉冲神经元前向传播
Args:
x: (B, T, in_features) 连续值输入
Returns:
spikes: (B, T, out_features) 二值脉冲序列
"""
B, T, _ = x.shape
weighted = self.weight(x)
membrane = torch.zeros(B, self.out_features, device=x.device)
spikes = torch.zeros(B, T, self.out_features, device=x.device)
for t in range(T):
membrane = membrane * (1 - 1/self.tau)
membrane = membrane + weighted[:, t, :]
fire = membrane > self.threshold
spikes[:, t, :] = fire.float()
membrane = membrane * (1 - fire.float())
return spikes
class FCASNN(nn.Module):
"""
Frequency and Channel Attention Spiking Neural Network
完整架构:
MVMD → IMF注意力 → 通道注意力 → SNN分类
"""
def __init__(self, num_channels: int = 14, num_imfs: int = 5,
hidden_dim: int = 128, num_classes: int = 2):
super().__init__()
self.mvmd = MVMD(num_channels, num_imfs)
self.imf_attn = IMFAttention(num_imfs)
self.channel_attn = ChannelAttention(num_channels)
self.flatten_dim = num_imfs * num_channels
self.sn1 = SpikingNeuronLayer(self.flatten_dim, hidden_dim)
self.sn2 = SpikingNeuronLayer(hidden_dim, hidden_dim // 2)
self.readout = nn.Sequential(
nn.Linear(hidden_dim // 2, hidden_dim // 4),
nn.ReLU(),
nn.Linear(hidden_dim // 4, num_classes)
)
def forward(self, eeg: torch.Tensor) -> dict:
"""
Args:
eeg: (B, C, T) 多通道EEG
Returns:
logits: (B, num_classes) 疲劳/清醒
spikes: 各层脉冲活动（用于分析）
"""
imfs = self.mvmd(eeg)
imfs_weighted = self.imf_attn(imfs)
B, M, C, T = imfs_weighted.shape
imf_channel_weighted = torch.zeros_like(imfs_weighted)
for m in range(M):
imf_channel_weighted[:, m] = self.channel_attn(
imfs_weighted[:, m]
)
features = imf_channel_weighted.permute(0, 3, 1, 2)
features = features.reshape(B, T, M * C)
s1 = self.sn1(features)
s2 = self.sn2(s1)
spike_rate = s2.mean(dim=1)
logits = self.readout(spike_rate)
return {
'logits': logits,
'imfs': imfs,
'imf_weights': self.imf_attn.scorer(
imfs.pow(2).mean(dim=(2, 3))
),
'spike_rate': spike_rate
}
if __name__ == "__main__":
B, C, T = 4, 14, 7500
np.random.seed(42)
t = np.linspace(0, 30, T)
eeg_alert = np.stack([
np.sin(2*np.pi*10*t) * 0.5 +
np.sin(2*np.pi*20*t) * 0.3 +
np.random.normal(0, 0.2, T)
for _ in range(C)
])
eeg_fatigue = np.stack([
np.sin(2*np.pi*5*t) * 0.8 +
np.sin(2*np.pi*10*t) * 0.3 +
np.random.normal(0, 0.2, T)
for _ in range(C)
])
eeg_data = torch.tensor(
np.stack([eeg_alert, eeg_alert, eeg_fatigue, eeg_fatigue]),
dtype=torch.float32
)
labels = torch.tensor([0, 0, 1, 1])
model = FCASNN(num_channels=C, num_imfs=5, hidden_dim=128)
with torch.no_grad():
output = model(eeg_data)
predictions = output['logits'].argmax(dim=-1)
print("=== MVMD-FCASNN 疲劳检测 ===")
print(f"输入: {eeg_data.shape} (B, C, T)")
print(f"IMFs: {output['imfs'].shape} (B, num_imfs, C, T)")
print(f"IMF注意力权重: {output['imf_weights'].mean(0).tolist()}")
print(f"脉冲发放率: {output['spike_rate'].mean(dim=0)[:4].tolist()}")
print(f"\n预测: {['清醒' if p==0 else '疲劳' for p in predictions.tolist()]}")
print(f"真实: {['清醒' if l==0 else '疲劳' for l in labels.tolist()]}")
print("\n=== 噪声鲁棒性测试 ===")
for snr_db in [20, 10, 0, -3, -6]:
noise_power = 1.0 / (10 ** (snr_db / 10))
noisy_eeg = eeg_data + torch.randn_like(eeg_data) * np.sqrt(noise_power)
with torch.no_grad():
noisy_out = model(noisy_eeg)
noisy_pred = noisy_out['logits'].argmax(dim=-1)
acc = (noisy_pred == labels).float().mean().item()
print(f" SNR={snr_db:>3}dB: 准确率={acc*100:.1f}%")
params = sum(p.numel() for p in model.parameters())
print(f"\n模型参数量: {params:,} ({params/1e6:.2f}M)")
print(f"约 {params/1e3:.0f}K 参数（轻量级）")
|