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
| """
DMS 时序状态追踪：循环网络 vs 前馈网络
核心问题：疲劳/分心是时间渐变过程
- 前馈网络：每帧独立判断，无时序上下文
- 循环网络：累积时序信息，追踪状态渐变
"""
import torch
import torch.nn as nn
import numpy as np
class FeedforwardDMS(nn.Module):
"""前馈 DMS：每帧独立"""
def __init__(self, feat_dim: int = 128, n_classes: int = 3):
super().__init__()
self.fc = nn.Sequential(
nn.Linear(feat_dim, 256),
nn.ReLU(),
nn.Linear(256, 128),
nn.ReLU(),
nn.Linear(128, n_classes),
)
def forward(self, x):
"""x: (B, T, F) → 每帧独立分类"""
return self.fc(x)
class RecurrentDMS(nn.Module):
"""循环 DMS：时序累积"""
def __init__(self, feat_dim: int = 128, hidden_dim: int = 256,
n_classes: int = 3, n_layers: int = 2):
super().__init__()
self.lstm = nn.LSTM(feat_dim, hidden_dim, n_layers,
batch_first=True, bidirectional=False)
self.classifier = nn.Linear(hidden_dim, n_classes)
def forward(self, x):
"""x: (B, T, F) → 累积时序"""
out, _ = self.lstm(x)
return self.classifier(out)
class MultiTimescaleDMS(nn.Module):
"""
多时间尺度 DMS
灵感来自 PLOS 论文：
- 短期循环（10s）：捕捉即时状态变化（眨眼/扫视）
- 中期循环（60s）：捕捉疲劳渐变（PERCLOS）
- 长期循环（5min）：捕捉深度疲劳（微睡眠模式）
"""
def __init__(self, feat_dim: int = 128, hidden_dim: int = 256, n_classes: int = 3):
super().__init__()
self.short_lstm = nn.LSTM(feat_dim, hidden_dim, 1, batch_first=True)
self.mid_lstm = nn.LSTM(hidden_dim, hidden_dim, 1, batch_first=True)
self.long_lstm = nn.LSTM(hidden_dim, hidden_dim, 1, batch_first=True)
self.classifier = nn.Linear(hidden_dim, n_classes)
def forward(self, x):
"""x: (B, T, F)"""
short_out, _ = self.short_lstm(x)
mid_out, _ = self.mid_lstm(short_out)
long_out, _ = self.long_lstm(mid_out)
return self.classifier(long_out)
if __name__ == "__main__":
ff = FeedforwardDMS()
rc = RecurrentDMS()
mt = MultiTimescaleDMS()
x = torch.randn(4, 300, 128)
print(f"前馈: {sum(p.numel() for p in ff.parameters()):,} params")
print(f"循环: {sum(p.numel() for p in rc.parameters()):,} params")
print(f"多尺度: {sum(p.numel() for p in mt.parameters()):,} params")
y_ff = ff(x)
y_rc = rc(x)
y_mt = mt(x)
print(f"\n输出形状: ff={y_ff.shape}, rc={y_rc.shape}, mt={y_mt.shape}")
|
