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
| import torch import torch.nn as nn
class EEGNet(nn.Module): """ EEG专用卷积网络 适用于实时分心检测 """ def __init__(self, n_channels=32, n_samples=256, n_classes=2): super().__init__() self.conv1 = nn.Conv2d(1, 16, (1, 64), padding=(0, 32)) self.bn1 = nn.BatchNorm2d(16) self.conv2 = nn.Conv2d(16, 32, (n_channels, 1), groups=16) self.bn2 = nn.BatchNorm2d(32) self.conv3 = nn.Conv2d(32, 64, (1, 16), padding=(0, 8)) self.bn3 = nn.BatchNorm2d(64) self.gap = nn.AdaptiveAvgPool2d((1, 1)) self.fc = nn.Linear(64, n_classes) def forward(self, x): """ 前向传播 Args: x: EEG信号 (B, 1, n_channels, n_samples) Returns: out: 分类结果 (B, n_classes) """ x = self.conv1(x) x = self.bn1(x) x = nn.functional.elu(x) x = self.conv2(x) x = self.bn2(x) x = nn.functional.elu(x) x = self.conv3(x) x = self.bn3(x) x = nn.functional.elu(x) x = self.gap(x) x = x.view(x.size(0), -1) out = self.fc(x) return out
class LSTMClassifier(nn.Module): """ LSTM时序分类器 适用于长时间序列分心检测 """ def __init__(self, input_size=32, hidden_size=64, num_layers=2, n_classes=2): super().__init__() self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, bidirectional=True) self.fc = nn.Linear(hidden_size * 2, n_classes) def forward(self, x): """ 前向传播 Args: x: EEG信号 (B, seq_len, input_size) Returns: out: 分类结果 (B, n_classes) """ lstm_out, _ = self.lstm(x) last_out = lstm_out[:, -1, :] out = self.fc(last_out) return out
if __name__ == "__main__": batch_size = 16 n_channels = 32 n_samples = 256 model = EEGNet(n_channels, n_samples, n_classes=2) eeg_input = torch.randn(batch_size, 1, n_channels, n_samples) with torch.no_grad(): output = model(eeg_input) print(f"EEGNet输出: {output.shape}") lstm_model = LSTMClassifier(input_size=n_channels) lstm_input = torch.randn(batch_size, 100, n_channels) with torch.no_grad(): lstm_output = lstm_model(lstm_input) print(f"LSTM输出: {lstm_output.shape}")
|