Buckets:
| """ | |
| model.py | |
| 1D CNN-LSTM architecture for 10-second multi-channel EEG seizure detection. | |
| Input shape: (batch_size, 18, 2560) | |
| Output shape: (batch_size, 2) [Logits for CrossEntropyLoss] | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| class ConvBlock(nn.Module): | |
| """ | |
| 1D Convolutional block comprising Conv1d, BatchNorm1d, ReLU, and MaxPool1d. | |
| Extracts local temporal-spatial patterns while reducing sequence length. | |
| """ | |
| def __init__(self, in_channels, out_channels, kernel_size=5, stride=2, pool_size=2): | |
| super().__init__() | |
| padding = kernel_size // 2 | |
| self.conv = nn.Conv1d( | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| kernel_size=kernel_size, | |
| stride=stride, | |
| padding=padding, | |
| bias=False | |
| ) | |
| self.bn = nn.BatchNorm1d(out_channels) | |
| self.act = nn.ReLU(inplace=True) | |
| self.pool = nn.MaxPool1d(kernel_size=pool_size, stride=pool_size) | |
| def forward(self, x): | |
| # x: (B, C_in, L_in) | |
| x = self.conv(x) # (B, C_out, L_in / stride) | |
| x = self.bn(x) | |
| x = self.act(x) | |
| x = self.pool(x) # (B, C_out, L_in / (stride * pool_size)) | |
| return x | |
| class CNNLSTM(nn.Module): | |
| """ | |
| CNN-LSTM baseline architecture for EEG Seizure Detection. | |
| Stages: | |
| 1. Multi-layer 1D CNN feature extractor to capture local multi-channel EEG morphology | |
| and downsample temporal dimension (2560 -> 40). | |
| 2. Temporal feature transposition: (B, C, T) -> (B, T, C). | |
| 3. Multi-layer Bidirectional LSTM to model long-range temporal dynamics across window. | |
| 4. Temporal pooling (Global Average Pooling) across LSTM hidden states. | |
| 5. Fully-connected classification head outputting 2 unnormalized logits. | |
| """ | |
| def __init__( | |
| self, | |
| in_channels=18, | |
| sequence_length=2560, | |
| conv_channels=None, | |
| conv_kernel_sizes=None, | |
| conv_strides=None, | |
| pool_kernel_size=2, | |
| pool_stride=2, | |
| lstm_hidden_size=128, | |
| lstm_num_layers=2, | |
| lstm_bidirectional=True, | |
| fc_hidden_size=64, | |
| dropout=0.3, | |
| lstm_dropout=None, | |
| classifier_dropout=None, | |
| feat_dropout=0.0, | |
| num_classes=2 | |
| ): | |
| super().__init__() | |
| if conv_channels is None: | |
| conv_channels = [32, 64, 128] | |
| if conv_kernel_sizes is None: | |
| conv_kernel_sizes = [7, 5, 3] | |
| if conv_strides is None: | |
| conv_strides = [2, 2, 2] | |
| # Regularization dropout rates | |
| self.lstm_dropout = dropout if lstm_dropout is None else lstm_dropout | |
| self.classifier_dropout = dropout if classifier_dropout is None else classifier_dropout | |
| self.feat_dropout_rate = feat_dropout | |
| self.in_channels = in_channels | |
| self.sequence_length = sequence_length | |
| self.lstm_hidden_size = lstm_hidden_size | |
| self.lstm_bidirectional = lstm_bidirectional | |
| # 1. Convolutional Stage | |
| conv_layers = [] | |
| curr_in = in_channels | |
| for out_ch, k, s in zip(conv_channels, conv_kernel_sizes, conv_strides): | |
| conv_layers.append(ConvBlock(curr_in, out_ch, kernel_size=k, stride=s, pool_size=pool_kernel_size)) | |
| curr_in = out_ch | |
| self.conv_net = nn.Sequential(*conv_layers) | |
| # Feature dropout between CNN and recurrent stage | |
| self.feat_dropout = nn.Dropout(p=feat_dropout) if feat_dropout > 0.0 else nn.Identity() | |
| # 2. LSTM Stage | |
| self.lstm = nn.LSTM( | |
| input_size=conv_channels[-1], | |
| hidden_size=lstm_hidden_size, | |
| num_layers=lstm_num_layers, | |
| batch_first=True, | |
| bidirectional=lstm_bidirectional, | |
| dropout=self.lstm_dropout if lstm_num_layers > 1 else 0.0 | |
| ) | |
| # 3. Dense Classifier Head | |
| lstm_out_dim = lstm_hidden_size * (2 if lstm_bidirectional else 1) | |
| self.classifier = nn.Sequential( | |
| nn.Dropout(p=self.classifier_dropout), | |
| nn.Linear(lstm_out_dim, fc_hidden_size), | |
| nn.ReLU(inplace=True), | |
| nn.Dropout(p=self.classifier_dropout), | |
| nn.Linear(fc_hidden_size, num_classes) | |
| ) | |
| def forward(self, x): | |
| """ | |
| Forward pass. | |
| Args: | |
| x: Tensor of shape (batch_size, 18, 2560) | |
| Returns: | |
| logits: Tensor of shape (batch_size, 2) | |
| """ | |
| # Stage 1: 1D CNN Feature Extraction | |
| # Input: (B, 18, 2560) | |
| feat = self.conv_net(x) | |
| # Output after 3 conv+pool blocks: (B, 128, 40) | |
| # Stage 2: Reshape into Temporal Sequence | |
| # Transpose from (B, channels, time) to (B, time, channels) | |
| feat_seq = feat.permute(0, 2, 1) # (B, 40, 128) | |
| feat_seq = self.feat_dropout(feat_seq) | |
| # Stage 3: Recurrent LSTM Modeling | |
| lstm_out, _ = self.lstm(feat_seq) # (B, 40, 256 if bidirectional else 128) | |
| # Stage 4: Temporal Aggregation (Global Average Pooling over time) | |
| pooled = torch.mean(lstm_out, dim=1) # (B, 256) | |
| # Stage 5: Classification Head | |
| logits = self.classifier(pooled) # (B, 2) | |
| return logits | |
| def build_model_from_config(config): | |
| """Factory function to build CNNLSTM model from config dictionary.""" | |
| model_cfg = config.get("model", {}) | |
| return CNNLSTM( | |
| in_channels=model_cfg.get("in_channels", 18), | |
| sequence_length=model_cfg.get("sequence_length", 2560), | |
| conv_channels=model_cfg.get("conv_channels", [32, 64, 128]), | |
| conv_kernel_sizes=model_cfg.get("conv_kernel_sizes", [7, 5, 3]), | |
| conv_strides=model_cfg.get("conv_strides", [2, 2, 2]), | |
| pool_kernel_size=model_cfg.get("pool_kernel_size", 2), | |
| pool_stride=model_cfg.get("pool_stride", 2), | |
| lstm_hidden_size=model_cfg.get("lstm_hidden_size", 128), | |
| lstm_num_layers=model_cfg.get("lstm_num_layers", 2), | |
| lstm_bidirectional=model_cfg.get("lstm_bidirectional", True), | |
| fc_hidden_size=model_cfg.get("fc_hidden_size", 64), | |
| dropout=model_cfg.get("dropout", 0.3), | |
| lstm_dropout=model_cfg.get("lstm_dropout", None), | |
| classifier_dropout=model_cfg.get("classifier_dropout", None), | |
| feat_dropout=model_cfg.get("feat_dropout", 0.0), | |
| num_classes=model_cfg.get("num_classes", 2) | |
| ) | |
| if __name__ == "__main__": | |
| # Smoke test for tensor shape verification | |
| model = CNNLSTM() | |
| dummy_input = torch.randn(4, 18, 2560) | |
| out = model(dummy_input) | |
| print("Model test successful!") | |
| print(f"Input shape: {dummy_input.shape}") | |
| print(f"Output logits shape: {out.shape}") | |
| total_params = sum(p.numel() for p in model.parameters() if p.requires_grad) | |
| print(f"Trainable parameters: {total_params:,}") | |
Xet Storage Details
- Size:
- 6.8 kB
- Xet hash:
- 501a9c1a2f30e4f7177296c4de278ef0f66ffa6abf5df1c3b3ae456e469a5705
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.