braindecode 0.8.1__py3-none-any.whl → 1.1.0__py3-none-any.whl

braindecode/models/deep4.py

@@ -3,17 +3,25 @@
 # License: BSD (3-clause)
 
 from einops.layers.torch import Rearrange
+from mne.utils import warn
 from torch import nn
 from torch.nn import init
-from torch.nn.functional import elu
 
-from .base import EEGModuleMixin, deprecated_args
-from .functions import identity, squeeze_final_output
-from .modules import AvgPool2dWithConv, CombinedConv, Ensure4d, Expression
+from braindecode.models.base import EEGModuleMixin
+from braindecode.modules import (
+    AvgPool2dWithConv,
+    CombinedConv,
+    Ensure4d,
+    SqueezeFinalOutput,
+)
 
 
 class Deep4Net(EEGModuleMixin, nn.Sequential):
-    """Deep ConvNet model from Schirrmeister et al 2017.
+    """Deep ConvNet model from Schirrmeister et al (2017) [Schirrmeister2017]_.
+
+     .. figure:: https://onlinelibrary.wiley.com/cms/asset/fc200ccc-d8c4-45b4-8577-56ce4d15999a/hbm23730-fig-0001-m.jpg
+        :align: center
+        :alt: CTNet Architecture
 
     Model described in [Schirrmeister2017]_.
 
@@ -44,13 +52,13 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
         Number of temporal filters in layer 4.
     filter_length_4: int
         Length of the temporal filter in layer 4.
-    first_conv_nonlin: callable
+    activation_first_conv_nonlin: nn.Module, default is nn.ELU
         Non-linear activation function to be used after convolution in layer 1.
     first_pool_mode: str
         Pooling mode in layer 1. "max" or "mean".
     first_pool_nonlin: callable
         Non-linear activation function to be used after pooling in layer 1.
-    later_conv_nonlin: callable
+    activation_later_conv_nonlin: nn.Module, default is nn.ELU
         Non-linear activation function to be used after convolution in later layers.
     later_pool_mode: str
         Pooling mode in later layers. "max" or "mean".
@@ -67,12 +75,6 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
         Momentum for BatchNorm2d.
     stride_before_pool: bool
         Stride before pooling.
-    in_chans :
-        Alias for n_chans.
-    n_classes:
-        Alias for n_outputs.
-    input_window_samples :
-        Alias for n_times.
 
 
     References
@@ -87,47 +89,38 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
     """
 
     def __init__(
-            self,
-            n_chans=None,
-            n_outputs=None,
-            n_times=None,
-            final_conv_length="auto",
-            n_filters_time=25,
-            n_filters_spat=25,
-            filter_time_length=10,
-            pool_time_length=3,
-            pool_time_stride=3,
-            n_filters_2=50,
-            filter_length_2=10,
-            n_filters_3=100,
-            filter_length_3=10,
-            n_filters_4=200,
-            filter_length_4=10,
-            first_conv_nonlin=elu,
-            first_pool_mode="max",
-            first_pool_nonlin=identity,
-            later_conv_nonlin=elu,
-            later_pool_mode="max",
-            later_pool_nonlin=identity,
-            drop_prob=0.5,
-            split_first_layer=True,
-            batch_norm=True,
-            batch_norm_alpha=0.1,
-            stride_before_pool=False,
-            chs_info=None,
-            input_window_seconds=None,
-            sfreq=None,
-            in_chans=None,
-            n_classes=None,
-            input_window_samples=None,
-            add_log_softmax=True,
+        self,
+        n_chans=None,
+        n_outputs=None,
+        n_times=None,
+        final_conv_length="auto",
+        n_filters_time=25,
+        n_filters_spat=25,
+        filter_time_length=10,
+        pool_time_length=3,
+        pool_time_stride=3,
+        n_filters_2=50,
+        filter_length_2=10,
+        n_filters_3=100,
+        filter_length_3=10,
+        n_filters_4=200,
+        filter_length_4=10,
+        activation_first_conv_nonlin: nn.Module = nn.ELU,
+        first_pool_mode="max",
+        first_pool_nonlin: nn.Module = nn.Identity,
+        activation_later_conv_nonlin: nn.Module = nn.ELU,
+        later_pool_mode="max",
+        later_pool_nonlin: nn.Module = nn.Identity,
+        drop_prob=0.5,
+        split_first_layer=True,
+        batch_norm=True,
+        batch_norm_alpha=0.1,
+        stride_before_pool=False,
+        # Braindecode EEGModuleMixin parameters
+        chs_info=None,
+        input_window_seconds=None,
+        sfreq=None,
     ):
-        n_chans, n_outputs, n_times = deprecated_args(
-            self,
-            ('in_chans', 'n_chans', in_chans, n_chans),
-            ('n_classes', 'n_outputs', n_classes, n_outputs),
-            ('input_window_samples', 'n_times', input_window_samples, n_times),
-        )
         super().__init__(
             n_outputs=n_outputs,
             n_chans=n_chans,
@@ -135,10 +128,9 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
             n_times=n_times,
             input_window_seconds=input_window_seconds,
             sfreq=sfreq,
-            add_log_softmax=add_log_softmax,
         )
         del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
-        del in_chans, n_classes, input_window_samples
+
         if final_conv_length == "auto":
             assert self.n_times is not None
         self.final_conv_length = final_conv_length
@@ -153,10 +145,10 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
         self.filter_length_3 = filter_length_3
         self.n_filters_4 = n_filters_4
         self.filter_length_4 = filter_length_4
-        self.first_nonlin = first_conv_nonlin
+        self.first_nonlin = activation_first_conv_nonlin
         self.first_pool_mode = first_pool_mode
         self.first_pool_nonlin = first_pool_nonlin
-        self.later_conv_nonlin = later_conv_nonlin
+        self.later_conv_nonlin = activation_later_conv_nonlin
         self.later_pool_mode = later_pool_mode
         self.later_pool_nonlin = later_pool_nonlin
         self.drop_prob = drop_prob
@@ -165,6 +157,27 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
         self.batch_norm_alpha = batch_norm_alpha
         self.stride_before_pool = stride_before_pool
 
+        min_n_times = self._get_min_n_times()
+        if self.n_times < min_n_times:
+            scaling_factor = self.n_times / min_n_times
+            warn(
+                f"n_times ({self.n_times}) is smaller than the minimum required "
+                f"({min_n_times}) for the current model parameters configuration. "
+                "Adjusting parameters to ensure compatibility."
+                "Reducing the kernel, pooling, and stride sizes accordingly."
+                "Scaling factor: {:.2f}".format(scaling_factor),
+                UserWarning,
+            )
+            # Calculate a scaling factor to adjust temporal parameters
+            # Apply the scaling factor to all temporal kernel and pooling sizes
+            self.filter_time_length = max(
+                1, int(self.filter_time_length * scaling_factor)
+            )
+            self.pool_time_length = max(1, int(self.pool_time_length * scaling_factor))
+            self.pool_time_stride = max(1, int(self.pool_time_stride * scaling_factor))
+            self.filter_length_2 = max(1, int(self.filter_length_2 * scaling_factor))
+            self.filter_length_3 = max(1, int(self.filter_length_3 * scaling_factor))
+            self.filter_length_4 = max(1, int(self.filter_length_4 * scaling_factor))
         # For the load_state_dict
         # When padronize all layers,
         # add the old's parameters here
@@ -174,7 +187,7 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
             "conv_time.bias": "conv_time_spat.conv_time.bias",
             "conv_spat.bias": "conv_time_spat.conv_spat.bias",
             "conv_classifier.weight": "final_layer.conv_classifier.weight",
-            "conv_classifier.bias": "final_layer.conv_classifier.bias"
+            "conv_classifier.bias": "final_layer.conv_classifier.bias",
         }
 
         if self.stride_before_pool:
@@ -223,17 +236,17 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
                     eps=1e-5,
                 ),
             )
-        self.add_module("conv_nonlin", Expression(self.first_nonlin))
+        self.add_module("conv_nonlin", self.first_nonlin())
         self.add_module(
             "pool",
             first_pool_class(
                 kernel_size=(self.pool_time_length, 1), stride=(pool_stride, 1)
             ),
         )
-        self.add_module("pool_nonlin", Expression(self.first_pool_nonlin))
+        self.add_module("pool_nonlin", self.first_pool_nonlin())
 
         def add_conv_pool_block(
-                model, n_filters_before, n_filters, filter_length, block_nr
+            model, n_filters_before, n_filters, filter_length, block_nr
         ):
             suffix = "_{:d}".format(block_nr)
             self.add_module("drop" + suffix, nn.Dropout(p=self.drop_prob))
@@ -257,7 +270,7 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
                         eps=1e-5,
                     ),
                 )
-            self.add_module("nonlin" + suffix, Expression(self.later_conv_nonlin))
+            self.add_module("nonlin" + suffix, self.later_conv_nonlin())
 
             self.add_module(
                 "pool" + suffix,
@@ -266,7 +279,7 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
                     stride=(pool_stride, 1),
                 ),
             )
-            self.add_module("pool_nonlin" + suffix, Expression(self.later_pool_nonlin))
+            self.add_module("pool_nonlin" + suffix, self.later_pool_nonlin())
 
         add_conv_pool_block(
             self, n_filters_conv, self.n_filters_2, self.filter_length_2, 2
@@ -278,7 +291,6 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
             self, self.n_filters_3, self.n_filters_4, self.filter_length_4, 4
         )
 
-        # self.add_module('drop_classifier', nn.Dropout(p=self.drop_prob))
         self.eval()
         if self.final_conv_length == "auto":
             self.final_conv_length = self.get_output_shape()[2]
@@ -286,17 +298,17 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
         # Incorporating classification module and subsequent ones in one final layer
         module = nn.Sequential()
 
-        module.add_module("conv_classifier",
-                          nn.Conv2d(
-                            self.n_filters_4,
-                            self.n_outputs,
-                            (self.final_conv_length, 1),
-                            bias=True, ))
-
-        if self.add_log_softmax:
-            module.add_module("logsoftmax", nn.LogSoftmax(dim=1))
+        module.add_module(
+            "conv_classifier",
+            nn.Conv2d(
+                self.n_filters_4,
+                self.n_outputs,
+                (self.final_conv_length, 1),
+                bias=True,
+            ),
+        )
 
-        module.add_module("squeeze", Expression(squeeze_final_output))
+        module.add_module("squeeze", SqueezeFinalOutput())
 
         self.add_module("final_layer", module)
 
@@ -309,7 +321,7 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
         if self.split_first_layer:
             init.xavier_uniform_(self.conv_time_spat.conv_spat.weight, gain=1)
             if not self.batch_norm:
-                init.constant_(self.conv_spat.bias, 0)
+                init.constant_(self.conv_time_spat.conv_spat.bias, 0)
         if self.batch_norm:
             init.constant_(self.bnorm.weight, 1)
             init.constant_(self.bnorm.bias, 0)
@@ -329,5 +341,32 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
         init.xavier_uniform_(self.final_layer.conv_classifier.weight, gain=1)
         init.constant_(self.final_layer.conv_classifier.bias, 0)
 
-        # Start in eval mode
-        self.eval()
+        self.train()
+
+    def _get_min_n_times(self) -> int:
+        """
+        Calculate the minimum number of time samples required for the model
+        to work with the given temporal parameters.
+        """
+        # Start with the minimum valid output length of the network (1)
+        min_len = 1
+
+        # List of conv kernel sizes and pool parameters for the 4 blocks, in reverse order
+        # Each tuple: (filter_length, pool_length, pool_stride)
+        block_params = [
+            (self.filter_length_4, self.pool_time_length, self.pool_time_stride),
+            (self.filter_length_3, self.pool_time_length, self.pool_time_stride),
+            (self.filter_length_2, self.pool_time_length, self.pool_time_stride),
+            (self.filter_time_length, self.pool_time_length, self.pool_time_stride),
+        ]
+
+        # Work backward from the last layer to the input
+        for filter_len, pool_len, pool_stride in block_params:
+            # Reverse the pooling operation
+            # L_in = stride * (L_out - 1) + kernel_size
+            min_len = pool_stride * (min_len - 1) + pool_len
+            # Reverse the convolution operation (assuming stride=1)
+            # L_in = L_out + kernel_size - 1
+            min_len = min_len + filter_len - 1
+
+        return min_len

braindecode/models/deepsleepnet.py

@@ -4,11 +4,133 @@
 import torch
 import torch.nn as nn
 
-from .base import EEGModuleMixin, deprecated_args
+from braindecode.models.base import EEGModuleMixin
 
 
-class _SmallCNN(nn.Module):  # smaller filter sizes to learn temporal information
-    def __init__(self):
+class DeepSleepNet(EEGModuleMixin, nn.Module):
+    """Sleep staging architecture from Supratak et al. (2017) [Supratak2017]_.
+
+     .. figure:: https://raw.githubusercontent.com/akaraspt/deepsleepnet/refs/heads/master/img/deepsleepnet.png
+        :align: center
+        :alt: DeepSleepNet Architecture
+
+    Convolutional neural network and bidirectional-Long Short-Term
+    for single channels sleep staging described in [Supratak2017]_.
+
+    Parameters
+    ----------
+    activation_large: nn.Module, default=nn.ELU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ELU``.
+    activation_small: nn.Module, default=nn.ReLU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ReLU``.
+    return_feats : bool
+        If True, return the features, i.e. the output of the feature extractor
+        (before the final linear layer). If False, pass the features through
+        the final linear layer.
+    drop_prob : float, default=0.5
+        The dropout rate for regularization. Values should be between 0 and 1.
+
+
+    References
+    ----------
+    .. [Supratak2017] Supratak, A., Dong, H., Wu, C., & Guo, Y. (2017).
+       DeepSleepNet: A model for automatic sleep stage scoring based
+       on raw single-channel EEG. IEEE Transactions on Neural Systems
+       and Rehabilitation Engineering, 25(11), 1998-2008.
+    """
+
+    def __init__(
+        self,
+        n_outputs=5,
+        return_feats=False,
+        n_chans=None,
+        chs_info=None,
+        n_times=None,
+        input_window_seconds=None,
+        sfreq=None,
+        activation_large: nn.Module = nn.ELU,
+        activation_small: nn.Module = nn.ReLU,
+        drop_prob: float = 0.5,
+    ):
+        super().__init__(
+            n_outputs=n_outputs,
+            n_chans=n_chans,
+            chs_info=chs_info,
+            n_times=n_times,
+            input_window_seconds=input_window_seconds,
+            sfreq=sfreq,
+        )
+        del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
+        self.cnn1 = _SmallCNN(activation=activation_small, drop_prob=drop_prob)
+        self.cnn2 = _LargeCNN(activation=activation_large, drop_prob=drop_prob)
+        self.dropout = nn.Dropout(0.5)
+        self.bilstm = _BiLSTM(input_size=3072, hidden_size=512, num_layers=2)
+        self.fc = nn.Sequential(
+            nn.Linear(3072, 1024, bias=False), nn.BatchNorm1d(num_features=1024)
+        )
+
+        self.features_extractor = nn.Identity()
+        self.len_last_layer = 1024
+        self.return_feats = return_feats
+
+        # TODO: Add new way to handle return_features == True
+        if not return_feats:
+            self.final_layer = nn.Linear(1024, self.n_outputs)
+        else:
+            self.final_layer = nn.Identity()
+
+    def forward(self, x):
+        """Forward pass.
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            Batch of EEG windows of shape (batch_size, n_channels, n_times).
+        """
+
+        if x.ndim == 3:
+            x = x.unsqueeze(1)
+
+        x1 = self.cnn1(x)
+        x1 = x1.flatten(start_dim=1)
+
+        x2 = self.cnn2(x)
+        x2 = x2.flatten(start_dim=1)
+
+        x = torch.cat((x1, x2), dim=1)
+        x = self.dropout(x)
+        temp = x.clone()
+        temp = self.fc(temp)
+        x = x.unsqueeze(1)
+        x = self.bilstm(x)
+        x = x.squeeze()
+        x = torch.add(x, temp)
+        x = self.dropout(x)
+
+        feats = self.features_extractor(x)
+
+        if self.return_feats:
+            return feats
+        else:
+            return self.final_layer(feats)
+
+
+class _SmallCNN(nn.Module):
+    """
+    Smaller filter sizes to learn temporal information.
+
+    Parameters
+    ----------
+    activation: nn.Module, default=nn.ReLU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ReLU``.
+    drop_prob : float, default=0.5
+        The dropout rate for regularization. Values should be between 0 and 1.
+    """
+
+    def __init__(self, activation: nn.Module = nn.ReLU, drop_prob: float = 0.5):
         super().__init__()
         self.conv1 = nn.Sequential(
             nn.Conv2d(
@@ -20,10 +142,10 @@ class _SmallCNN(nn.Module):  # smaller filter sizes to learn temporal informatio
                 bias=False,
             ),
             nn.BatchNorm2d(num_features=64),
-            nn.ReLU(),
+            activation(),
         )
         self.pool1 = nn.MaxPool2d(kernel_size=(1, 8), stride=(1, 8), padding=(0, 2))
-        self.dropout = nn.Dropout(p=0.5)
+        self.dropout = nn.Dropout(p=drop_prob)
         self.conv2 = nn.Sequential(
             nn.Conv2d(
                 in_channels=64,
@@ -34,7 +156,7 @@ class _SmallCNN(nn.Module):  # smaller filter sizes to learn temporal informatio
                 bias=False,
             ),
             nn.BatchNorm2d(num_features=128),
-            nn.ReLU(),
+            activation(),
         )
         self.conv3 = nn.Sequential(
             nn.Conv2d(
@@ -46,7 +168,7 @@ class _SmallCNN(nn.Module):  # smaller filter sizes to learn temporal informatio
                 bias=False,
             ),
             nn.BatchNorm2d(num_features=128),
-            nn.ReLU(),
+            activation(),
         )
         self.conv4 = nn.Sequential(
             nn.Conv2d(
@@ -58,7 +180,7 @@ class _SmallCNN(nn.Module):  # smaller filter sizes to learn temporal informatio
                 bias=False,
             ),
             nn.BatchNorm2d(num_features=128),
-            nn.ReLU(),
+            activation(),
         )
         self.pool2 = nn.MaxPool2d(kernel_size=(1, 4), stride=(1, 4), padding=(0, 1))
 
@@ -72,8 +194,19 @@ class _SmallCNN(nn.Module):  # smaller filter sizes to learn temporal informatio
         return x
 
 
-class _LargeCNN(nn.Module):  # larger filter sizes to learn frequency information
-    def __init__(self):
+class _LargeCNN(nn.Module):
+    """
+    Larger filter sizes to learn frequency information.
+
+    Parameters
+    ----------
+    activation: nn.Module, default=nn.ELU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ELU``.
+
+    """
+
+    def __init__(self, activation: nn.Module = nn.ELU, drop_prob: float = 0.5):
         super().__init__()
 
         self.conv1 = nn.Sequential(
@@ -86,10 +219,10 @@ class _LargeCNN(nn.Module):  # larger filter sizes to learn frequency informatio
                 bias=False,
             ),
             nn.BatchNorm2d(num_features=64),
-            nn.ReLU(),
+            activation(),
         )
         self.pool1 = nn.MaxPool2d(kernel_size=(1, 4), stride=(1, 4))
-        self.dropout = nn.Dropout(p=0.5)
+        self.dropout = nn.Dropout(p=drop_prob)
         self.conv2 = nn.Sequential(
             nn.Conv2d(
                 in_channels=64,
@@ -100,7 +233,7 @@ class _LargeCNN(nn.Module):  # larger filter sizes to learn frequency informatio
                 bias=False,
             ),
             nn.BatchNorm2d(num_features=128),
-            nn.ReLU(),
+            activation(),
         )
         self.conv3 = nn.Sequential(
             nn.Conv2d(
@@ -112,7 +245,7 @@ class _LargeCNN(nn.Module):  # larger filter sizes to learn frequency informatio
                 bias=False,
             ),
             nn.BatchNorm2d(num_features=128),
-            nn.ReLU(),
+            activation(),
         )
         self.conv4 = nn.Sequential(
             nn.Conv2d(
@@ -124,7 +257,7 @@ class _LargeCNN(nn.Module):  # larger filter sizes to learn frequency informatio
                 bias=False,
             ),
             nn.BatchNorm2d(num_features=128),
-            nn.ReLU(),
+            activation(),
         )
         self.pool2 = nn.MaxPool2d(kernel_size=(1, 2), stride=(1, 2), padding=(0, 1))
 
@@ -154,112 +287,9 @@ class _BiLSTM(nn.Module):
 
     def forward(self, x):
         # set initial hidden and cell states
-        h0 = torch.zeros(
-            self.num_layers * 2, x.size(0), self.hidden_size
-        ).to(x.device)
+        h0 = torch.zeros(self.num_layers * 2, x.size(0), self.hidden_size).to(x.device)
         c0 = torch.zeros(self.num_layers * 2, x.size(0), self.hidden_size).to(x.device)
 
         # forward propagate LSTM
         out, _ = self.lstm(x, (h0, c0))
         return out
-
-
-class DeepSleepNet(EEGModuleMixin, nn.Module):
-    """Sleep staging architecture from Supratak et al 2017.
-
-    Convolutional neural network and bidirectional-Long Short-Term
-    for single channels sleep staging described in [Supratak2017]_.
-
-    Parameters
-    ----------
-    return_feats : bool
-        If True, return the features, i.e. the output of the feature extractor
-        (before the final linear layer). If False, pass the features through
-        the final linear layer.
-    n_classes :
-        Alias for n_outputs.
-
-    References
-    ----------
-    .. [Supratak2017] Supratak, A., Dong, H., Wu, C., & Guo, Y. (2017).
-       DeepSleepNet: A model for automatic sleep stage scoring based
-       on raw single-channel EEG. IEEE Transactions on Neural Systems
-       and Rehabilitation Engineering, 25(11), 1998-2008.
-    """
-
-    def __init__(
-            self,
-            n_outputs=5,
-            return_feats=False,
-            n_chans=None,
-            chs_info=None,
-            n_times=None,
-            input_window_seconds=None,
-            sfreq=None,
-            n_classes=None,
-    ):
-        n_outputs, = deprecated_args(
-            self,
-            ('n_classes', 'n_outputs', n_classes, n_outputs),
-        )
-        super().__init__(
-            n_outputs=n_outputs,
-            n_chans=n_chans,
-            chs_info=chs_info,
-            n_times=n_times,
-            input_window_seconds=input_window_seconds,
-            sfreq=sfreq,
-        )
-        del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
-        del n_classes
-        self.cnn1 = _SmallCNN()
-        self.cnn2 = _LargeCNN()
-        self.dropout = nn.Dropout(0.5)
-        self.bilstm = _BiLSTM(input_size=3072, hidden_size=512, num_layers=2)
-        self.fc = nn.Sequential(nn.Linear(3072, 1024, bias=False),
-                                nn.BatchNorm1d(num_features=1024))
-
-        self.features_extractor = nn.Identity()
-        self.len_last_layer = 1024
-        self.return_feats = return_feats
-
-        # TODO: Add new way to handle return_features == True
-        if not return_feats:
-            self.final_layer = nn.Linear(1024, self.n_outputs)
-        else:
-            self.final_layer = nn.Identity()
-
-    def forward(self, x):
-        """Forward pass.
-
-        Parameters
-        ----------
-        x: torch.Tensor
-            Batch of EEG windows of shape (batch_size, n_channels, n_times).
-        """
-
-        if x.ndim == 3:
-            x = x.unsqueeze(1)
-
-        x1 = self.cnn1(x)
-        x1 = x1.flatten(start_dim=1)
-
-        x2 = self.cnn2(x)
-        x2 = x2.flatten(start_dim=1)
-
-        x = torch.cat((x1, x2), dim=1)
-        x = self.dropout(x)
-        temp = x.clone()
-        temp = self.fc(temp)
-        x = x.unsqueeze(1)
-        x = self.bilstm(x)
-        x = x.squeeze()
-        x = torch.add(x, temp)
-        x = self.dropout(x)
-
-        feats = self.features_extractor(x)
-
-        if self.return_feats:
-            return feats
-        else:
-            return self.final_layer(feats)