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

braindecode/models/sinc_shallow.py

@@ -0,0 +1,337 @@
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops.layers.torch import Rearrange
+
+from braindecode.models.base import EEGModuleMixin
+
+
+class SincShallowNet(EEGModuleMixin, nn.Module):
+    """Sinc-ShallowNet from Borra, D et al (2020) [borra2020]_.
+
+    .. figure:: https://ars.els-cdn.com/content/image/1-s2.0-S0893608020302021-gr2_lrg.jpg
+        :align: center
+        :alt: SincShallowNet Architecture
+
+    The Sinc-ShallowNet architecture has these fundamental blocks:
+
+    1. **Block 1: Spectral and Spatial Feature Extraction**
+        - *Temporal Sinc-Convolutional Layer*:
+            Uses parametrized sinc functions to learn band-pass filters,
+            significantly reducing the number of trainable parameters by only
+            learning the lower and upper cutoff frequencies for each filter.
+       - *Spatial Depthwise Convolutional Layer*:
+            Applies depthwise convolutions to learn spatial filters for
+            each temporal feature map independently, further reducing
+            parameters and enhancing interpretability.
+       - *Batch Normalization*
+
+    2. **Block 2: Temporal Aggregation**
+        - *Activation Function*: ELU
+        - *Average Pooling Layer*: Aggregation by averaging spatial dim
+        - *Dropout Layer*
+        - *Flatten Layer*
+
+    3. **Block 3: Classification**
+        - *Fully Connected Layer*: Maps the feature vector to n_outputs.
+
+    **Implementation Notes:**
+
+    - The sinc-convolutional layer initializes cutoff frequencies uniformly
+        within the desired frequency range and updates them during training while
+        ensuring the lower cutoff is less than the upper cutoff.
+
+    Parameters
+    ----------
+    num_time_filters : int
+        Number of temporal filters in the SincFilter layer.
+    time_filter_len : int
+        Size of the temporal filters.
+    depth_multiplier : int
+        Depth multiplier for spatial filtering.
+    activation : nn.Module, optional
+        Activation function to use. Default is nn.ELU().
+    drop_prob : float, optional
+        Dropout probability. Default is 0.5.
+    first_freq : float, optional
+        The starting frequency for the first Sinc filter. Default is 5.0.
+    min_freq : float, optional
+        Minimum frequency allowed for the low frequencies of the filters. Default is 1.0.
+    freq_stride : float, optional
+        Frequency stride for the Sinc filters. Controls the spacing between the filter frequencies.
+        Default is 1.0.
+    padding : str, optional
+        Padding mode for convolution, either 'same' or 'valid'. Default is 'same'.
+    bandwidth : float, optional
+        Initial bandwidth for each Sinc filter. Default is 4.0.
+    pool_size : int, optional
+        Size of the pooling window for the average pooling layer. Default is 55.
+    pool_stride : int, optional
+        Stride of the pooling operation. Default is 12.
+
+    Notes
+    -----
+    This implementation is based on the implementation from [sincshallowcode]_.
+
+    References
+    ----------
+    .. [borra2020] Borra, D., Fantozzi, S., & Magosso, E. (2020). Interpretable
+       and lightweight convolutional neural network for EEG decoding: Application
+       to movement execution and imagination. Neural Networks, 129, 55-74.
+    .. [sincshallowcode] Sinc-ShallowNet re-implementation source code:
+       https://github.com/marcellosicbaldi/SincNet-Tensorflow
+    """
+
+    def __init__(
+        self,
+        num_time_filters: int = 32,
+        time_filter_len: int = 33,
+        depth_multiplier: int = 2,
+        activation: Optional[nn.Module] = nn.ELU,
+        drop_prob: float = 0.5,
+        first_freq: float = 5.0,
+        min_freq: float = 1.0,
+        freq_stride: float = 1.0,
+        padding: str = "same",
+        bandwidth: float = 4.0,
+        pool_size: int = 55,
+        pool_stride: int = 12,
+        # braindecode parameters
+        n_chans=None,
+        n_outputs=None,
+        n_times=None,
+        input_window_seconds=None,
+        sfreq=None,
+        chs_info=None,
+    ):
+        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
+
+        # Checkers and creating variables
+        if activation is None:
+            activation = nn.ELU()
+
+        # Define low frequencies for the SincFilter
+        low_freqs = torch.arange(
+            first_freq,
+            first_freq + num_time_filters * freq_stride,
+            freq_stride,
+            dtype=torch.float32,
+        )
+        self.n_filters = len(low_freqs)
+
+        if padding.lower() == "valid":
+            n_times_after_sinc_filter = self.n_times - time_filter_len + 1
+        elif padding.lower() == "same":
+            n_times_after_sinc_filter = self.n_times
+        else:
+            raise ValueError("Padding must be 'valid' or 'same'.")
+
+        size_after_pooling = (
+            (n_times_after_sinc_filter - pool_size) // pool_stride
+        ) + 1
+        flattened_size = num_time_filters * depth_multiplier * size_after_pooling
+
+        # Layers
+        self.ensuredims = Rearrange("batch chans times -> batch chans times 1")
+
+        # Block 1: Sinc filter
+        self.sinc_filter_layer = _SincFilter(
+            low_freqs=low_freqs,
+            kernel_size=time_filter_len,
+            sfreq=self.sfreq,
+            padding=padding,
+            bandwidth=bandwidth,
+            min_freq=min_freq,
+        )
+
+        self.depthwiseconv = nn.Sequential(
+            # Matching dim to depth wise conv!
+            Rearrange("batch timefil time nfilter -> batch nfilter timefil time"),
+            nn.BatchNorm2d(
+                self.n_filters, momentum=0.99
+            ),  # To match keras implementation
+            nn.Conv2d(
+                in_channels=self.n_filters,
+                out_channels=depth_multiplier * self.n_filters,
+                kernel_size=(self.n_chans, 1),
+                groups=self.n_filters,
+                bias=False,
+            ),
+        )
+
+        # Block 2: Batch norm, activation, pooling, dropout
+        self.temporal_aggregation = nn.Sequential(
+            nn.BatchNorm2d(depth_multiplier * self.n_filters, momentum=0.99),
+            activation(),
+            nn.AvgPool2d(kernel_size=(1, pool_size), stride=(1, pool_stride)),
+            nn.Dropout(p=drop_prob),
+            nn.Flatten(),
+        )
+
+        # Final classification layer
+        self.final_layer = nn.Linear(
+            flattened_size,
+            self.n_outputs,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the model.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape [batch_size, num_channels, num_samples].
+
+        Returns
+        -------
+        torch.Tensor
+            Output logits of shape [batch_size, num_classes].
+        """
+        x = self.ensuredims(x)
+        x = self.sinc_filter_layer(x)
+        x = self.depthwiseconv(x)
+        x = self.temporal_aggregation(x)
+
+        return self.final_layer(x)
+
+
+class _SincFilter(nn.Module):
+    """Sinc-Based Convolutional Layer for Band-Pass Filtering from Ravanelli and Bengio (2018) [ravanelli]_.
+
+    The `SincFilter` layer implements a convolutional layer where each kernel is
+    defined using a parametrized sinc function.
+    This design enforces each kernel to represent a band-pass filter,
+    reducing the number of trainable parameters.
+
+    Parameters
+    ----------
+    low_freqs : torch.Tensor
+        Initial low cutoff frequencies for each filter.
+    kernel_size : int
+        Size of the convolutional kernels (filters). Must be odd.
+    sfreq : float
+        Sampling rate of the input signal.
+    bandwidth : float, optional
+        Initial bandwidth for each filter. Default is 4.0.
+    min_freq : float, optional
+        Minimum frequency allowed for low frequencies. Default is 1.0.
+    padding : str, optional
+        Padding mode, either 'same' or 'valid'. Default is 'same'.
+
+    References
+    ----------
+    .. [ravanelli] Ravanelli, M., & Bengio, Y. (2018, December). Speaker
+       recognition from raw waveform with sincnet. In 2018 IEEE spoken language
+       technology workshop (SLT) (pp. 1021-1028). IEEE.
+    """
+
+    def __init__(
+        self,
+        low_freqs: torch.Tensor,
+        kernel_size: int,
+        sfreq: float,
+        bandwidth: float = 4.0,
+        min_freq: float = 1.0,
+        padding: str = "same",
+    ):
+        super().__init__()
+        if kernel_size % 2 == 0:
+            raise ValueError("Kernel size must be odd.")
+
+        self.num_filters = low_freqs.numel()
+        self.kernel_size = kernel_size
+        self.sfreq = sfreq
+        self.min_freq = min_freq
+        self.padding = padding.lower()
+
+        # Precompute constants
+        window = torch.hamming_window(kernel_size, periodic=False)
+
+        self.register_buffer("window", window[: kernel_size // 2].unsqueeze(-1))
+
+        n_pi = (
+            torch.arange(-(kernel_size // 2), 0, dtype=torch.float32)
+            / sfreq
+            * 2
+            * math.pi
+        )
+        self.register_buffer("n_pi", n_pi.unsqueeze(-1))
+
+        # Initialize learnable parameters
+        bandwidths = torch.full((1, self.num_filters), bandwidth)
+        self.bandwidths = nn.Parameter(bandwidths)
+        self.low_freqs = nn.Parameter(low_freqs.unsqueeze(0))
+
+        # Constant tensor of ones for filter construction
+        self.register_buffer("ones", torch.ones(1, 1, 1, self.num_filters))
+
+    def build_sinc_filters(self) -> torch.Tensor:
+        """Builds the sinc filters based on current parameters."""
+        # Computing the low frequencies of the filters
+        low_freqs = self.min_freq + torch.abs(self.low_freqs)
+        # Setting a minimum band and minimum freq
+        high_freqs = torch.clamp(
+            low_freqs + torch.abs(self.bandwidths),
+            min=self.min_freq,
+            max=self.sfreq / 2.0,
+        )
+        bandwidths = high_freqs - low_freqs
+
+        # Passing from n_ to the corresponding f_times_t domain
+        low = self.n_pi * low_freqs  # [kernel_size // 2, num_filters]
+        high = self.n_pi * high_freqs  # [kernel_size // 2, num_filters]
+
+        filters_left = (torch.sin(high) - torch.sin(low)) / (self.n_pi / 2.0)
+        filters_left *= self.window
+        filters_left /= 2.0 * bandwidths
+
+        # [1, kernel_size // 2, 1, num_filters]
+        filters_left = filters_left.unsqueeze(0).unsqueeze(2)
+        filters_right = torch.flip(filters_left, dims=[1])
+
+        filters = torch.cat(
+            [filters_left, self.ones, filters_right], dim=1
+        )  # [1, kernel_size, 1, num_filters]
+        filters = filters / torch.std(filters)
+        return filters
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        """
+        Apply sinc filters to the input signal.
+
+        Parameters
+        ----------
+        inputs : torch.Tensor
+            Input tensor of shape [batch_size, num_channels, num_samples, 1].
+
+        Returns
+        -------
+        torch.Tensor
+            Filtered output tensor of shape [batch_size, num_channels, num_samples, num_filters].
+        """
+        filters = self.build_sinc_filters().to(
+            inputs.device
+        )  # [1, kernel_size, 1, num_filters]
+
+        # Convert from channels_last to channels_first format
+        inputs = inputs.permute(0, 3, 1, 2)
+        # Permuting to match conv:
+        filters = filters.permute(3, 2, 0, 1)
+        # Apply convolution
+        outputs = F.conv2d(inputs, filters, padding=self.padding)
+        # Changing the dimensional
+        outputs = outputs.permute(0, 2, 3, 1)
+
+        return outputs

braindecode/models/sleep_stager_blanco_2020.py

@@ -5,11 +5,15 @@
 import torch
 from torch import nn
 
-from .base import EEGModuleMixin, deprecated_args
+from braindecode.models.base import EEGModuleMixin
 
 
 class SleepStagerBlanco2020(EEGModuleMixin, nn.Module):
-    """Sleep staging architecture from Blanco et al 2020.
+    """Sleep staging architecture from Blanco et al. (2020) from [Blanco2020]_
+
+    .. figure:: https://media.springernature.com/full/springer-static/image/art%3A10.1007%2Fs00500-019-04174-1/MediaObjects/500_2019_4174_Fig2_HTML.png
+        :align: center
+        :alt: SleepStagerBlanco2020 Architecture
 
     Convolutional neural network for sleep staging described in [Blanco2020]_.
     A series of seven convolutional layers with kernel sizes running down from 7 to 3,
@@ -24,7 +28,7 @@ class SleepStagerBlanco2020(EEGModuleMixin, nn.Module):
         Number of groups for the convolution. Set to 2 in [Blanco2020]_ for 2 Channel EEG.
         controls the connections between inputs and outputs. n_channels and n_conv_chans must be
         divisible by n_groups.
-    dropout : float
+    drop_prob : float
         Dropout rate before the output dense layer.
     apply_batch_norm : bool
         If True, apply batch normalization after both temporal convolutional
@@ -39,6 +43,9 @@ class SleepStagerBlanco2020(EEGModuleMixin, nn.Module):
         Alias for `n_outputs`.
     input_size_s : float
         Alias for `input_window_seconds`.
+    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``.
 
     References
     ----------
@@ -48,30 +55,21 @@ class SleepStagerBlanco2020(EEGModuleMixin, nn.Module):
     """
 
     def __init__(
-            self,
-            n_chans=None,
-            sfreq=None,
-            n_conv_chans=20,
-            input_window_seconds=30,
-            n_outputs=5,
-            n_groups=2,
-            max_pool_size=2,
-            dropout=0.5,
-            apply_batch_norm=False,
-            return_feats=False,
-            chs_info=None,
-            n_times=None,
-            n_channels=None,
-            n_classes=None,
-            input_size_s=None,
-            add_log_softmax=True,
+        self,
+        n_chans=None,
+        sfreq=None,
+        n_conv_chans=20,
+        input_window_seconds=None,
+        n_outputs=5,
+        n_groups=2,
+        max_pool_size=2,
+        drop_prob=0.5,
+        apply_batch_norm=False,
+        return_feats=False,
+        activation: nn.Module = nn.ReLU,
+        chs_info=None,
+        n_times=None,
     ):
-        n_chans, n_outputs, input_window_seconds, = deprecated_args(
-            self,
-            ("n_channels", "n_chans", n_channels, n_chans),
-            ("n_classes", "n_outputs", n_classes, n_outputs),
-            ("input_size_s", "input_window_seconds", input_size_s, input_window_seconds),
-        )
         super().__init__(
             n_outputs=n_outputs,
             n_chans=n_chans,
@@ -79,14 +77,12 @@ class SleepStagerBlanco2020(EEGModuleMixin, nn.Module):
             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 n_channels, n_classes, input_size_s
 
         self.mapping = {
             "fc.1.weight": "final_layer.1.weight",
-            "fc.1.bias": "final_layer.1.bias"
+            "fc.1.bias": "final_layer.1.bias",
         }
 
         batch_norm = nn.BatchNorm2d if apply_batch_norm else nn.Identity
@@ -94,32 +90,44 @@ class SleepStagerBlanco2020(EEGModuleMixin, nn.Module):
         self.feature_extractor = nn.Sequential(
             nn.Conv2d(self.n_chans, n_conv_chans, (1, 7), groups=n_groups, padding=0),
             batch_norm(n_conv_chans),
-            nn.ReLU(),
+            activation(),
             nn.MaxPool2d((1, max_pool_size)),
-            nn.Conv2d(n_conv_chans, n_conv_chans, (1, 7), groups=n_conv_chans, padding=0),
+            nn.Conv2d(
+                n_conv_chans, n_conv_chans, (1, 7), groups=n_conv_chans, padding=0
+            ),
             batch_norm(n_conv_chans),
-            nn.ReLU(),
+            activation(),
             nn.MaxPool2d((1, max_pool_size)),
-            nn.Conv2d(n_conv_chans, n_conv_chans, (1, 5), groups=n_conv_chans, padding=0),
+            nn.Conv2d(
+                n_conv_chans, n_conv_chans, (1, 5), groups=n_conv_chans, padding=0
+            ),
             batch_norm(n_conv_chans),
-            nn.ReLU(),
+            activation(),
             nn.MaxPool2d((1, max_pool_size)),
-            nn.Conv2d(n_conv_chans, n_conv_chans, (1, 5), groups=n_conv_chans, padding=0),
+            nn.Conv2d(
+                n_conv_chans, n_conv_chans, (1, 5), groups=n_conv_chans, padding=0
+            ),
             batch_norm(n_conv_chans),
-            nn.ReLU(),
+            activation(),
             nn.MaxPool2d((1, max_pool_size)),
-            nn.Conv2d(n_conv_chans, n_conv_chans, (1, 5), groups=n_conv_chans, padding=0),
+            nn.Conv2d(
+                n_conv_chans, n_conv_chans, (1, 5), groups=n_conv_chans, padding=0
+            ),
             batch_norm(n_conv_chans),
-            nn.ReLU(),
+            activation(),
             nn.MaxPool2d((1, max_pool_size)),
-            nn.Conv2d(n_conv_chans, n_conv_chans, (1, 3), groups=n_conv_chans, padding=0),
+            nn.Conv2d(
+                n_conv_chans, n_conv_chans, (1, 3), groups=n_conv_chans, padding=0
+            ),
             batch_norm(n_conv_chans),
-            nn.ReLU(),
+            activation(),
             nn.MaxPool2d((1, max_pool_size)),
-            nn.Conv2d(n_conv_chans, n_conv_chans, (1, 3), groups=n_conv_chans, padding=0),
+            nn.Conv2d(
+                n_conv_chans, n_conv_chans, (1, 3), groups=n_conv_chans, padding=0
+            ),
             batch_norm(n_conv_chans),
-            nn.ReLU(),
-            nn.MaxPool2d((1, max_pool_size))
+            activation(),
+            nn.MaxPool2d((1, max_pool_size)),
         )
 
         self.len_last_layer = self._len_last_layer(self.n_chans, self.n_times)
@@ -128,16 +136,17 @@ class SleepStagerBlanco2020(EEGModuleMixin, nn.Module):
         # TODO: Add new way to handle return_features == True
         if not return_feats:
             self.final_layer = nn.Sequential(
-                nn.Dropout(dropout),
+                nn.Dropout(drop_prob),
                 nn.Linear(self.len_last_layer, self.n_outputs),
-                nn.LogSoftmax(dim=1) if self.add_log_softmax else nn.Identity()
+                nn.Identity(),
             )
 
     def _len_last_layer(self, n_channels, input_size):
         self.feature_extractor.eval()
         with torch.no_grad():
             out = self.feature_extractor(
-                torch.Tensor(1, n_channels, 1, input_size))  # batch_size,n_channels,height,width
+                torch.Tensor(1, n_channels, 1, input_size)
+            )  # batch_size,n_channels,height,width
         self.feature_extractor.train()
         return len(out.flatten())
 

braindecode/models/sleep_stager_chambon_2018.py

@@ -2,14 +2,20 @@
 #
 # License: BSD (3-clause)
 
+import math
+
 import torch
 from torch import nn
-import numpy as np
-from .base import EEGModuleMixin, deprecated_args
+
+from braindecode.models.base import EEGModuleMixin
 
 
 class SleepStagerChambon2018(EEGModuleMixin, nn.Module):
-    """Sleep staging architecture from Chambon et al 2018.
+    """Sleep staging architecture from Chambon et al. (2018) [Chambon2018]_.
+
+    .. figure:: https://braindecode.org/dev/_static/model/SleepStagerChambon2018.jpg
+        :align: center
+        :alt: SleepStagerChambon2018 Architecture
 
     Convolutional neural network for sleep staging described in [Chambon2018]_.
 
@@ -26,7 +32,7 @@ class SleepStagerChambon2018(EEGModuleMixin, nn.Module):
     pad_size_s : float
         Padding size, in seconds. Set to 0.25 in [Chambon2018]_ (half the
         temporal convolution kernel size).
-    dropout : float
+    drop_prob : float
         Dropout rate before the output dense layer.
     apply_batch_norm : bool
         If True, apply batch normalization after both temporal convolutional
@@ -41,6 +47,9 @@ class SleepStagerChambon2018(EEGModuleMixin, nn.Module):
         Alias for `input_window_seconds`.
     n_classes:
         Alias for `n_outputs`.
+    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``.
 
     References
     ----------
@@ -52,30 +61,22 @@ class SleepStagerChambon2018(EEGModuleMixin, nn.Module):
     """
 
     def __init__(
-            self,
-            n_chans=None,
-            sfreq=None,
-            n_conv_chs=8,
-            time_conv_size_s=0.5,
-            max_pool_size_s=0.125,
-            pad_size_s=0.25,
-            input_window_seconds=30,
-            n_outputs=5,
-            dropout=0.25,
-            apply_batch_norm=False,
-            return_feats=False,
-            chs_info=None,
-            n_times=None,
-            n_channels=None,
-            input_size_s=None,
-            n_classes=None,
+        self,
+        n_chans=None,
+        sfreq=None,
+        n_conv_chs=8,
+        time_conv_size_s=0.5,
+        max_pool_size_s=0.125,
+        pad_size_s=0.25,
+        activation: nn.Module = nn.ReLU,
+        input_window_seconds=None,
+        n_outputs=5,
+        drop_prob=0.25,
+        apply_batch_norm=False,
+        return_feats=False,
+        chs_info=None,
+        n_times=None,
     ):
-        n_chans, n_outputs, input_window_seconds, = deprecated_args(
-            self,
-            ("n_channels", "n_chans", n_channels, n_chans),
-            ("n_classes", "n_outputs", n_classes, n_outputs),
-            ("input_size_s", "input_window_seconds", input_size_s, input_window_seconds),
-        )
         super().__init__(
             n_outputs=n_outputs,
             n_chans=n_chans,
@@ -85,54 +86,54 @@ class SleepStagerChambon2018(EEGModuleMixin, nn.Module):
             sfreq=sfreq,
         )
         del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
-        del n_channels, n_classes, input_size_s
 
         self.mapping = {
             "fc.1.weight": "final_layer.1.weight",
-            "fc.1.bias": "final_layer.1.bias"
+            "fc.1.bias": "final_layer.1.bias",
         }
 
-        time_conv_size = np.ceil(time_conv_size_s * self.sfreq).astype(int)
-        max_pool_size = np.ceil(max_pool_size_s * self.sfreq).astype(int)
-        pad_size = np.ceil(pad_size_s * self.sfreq).astype(int)
+        time_conv_size = math.ceil(time_conv_size_s * self.sfreq)
+        max_pool_size = math.ceil(max_pool_size_s * self.sfreq)
+        pad_size = math.ceil(pad_size_s * self.sfreq)
 
         if self.n_chans > 1:
             self.spatial_conv = nn.Conv2d(1, self.n_chans, (self.n_chans, 1))
+        else:
+            self.spatial_conv = nn.Identity()
 
         batch_norm = nn.BatchNorm2d if apply_batch_norm else nn.Identity
 
         self.feature_extractor = nn.Sequential(
-            nn.Conv2d(
-                1, n_conv_chs, (1, time_conv_size), padding=(0, pad_size)),
+            nn.Conv2d(1, n_conv_chs, (1, time_conv_size), padding=(0, pad_size)),
             batch_norm(n_conv_chs),
-            nn.ReLU(),
+            activation(),
             nn.MaxPool2d((1, max_pool_size)),
             nn.Conv2d(
-                n_conv_chs, n_conv_chs, (1, time_conv_size),
-                padding=(0, pad_size)),
+                n_conv_chs, n_conv_chs, (1, time_conv_size), padding=(0, pad_size)
+            ),
             batch_norm(n_conv_chs),
-            nn.ReLU(),
-            nn.MaxPool2d((1, max_pool_size))
+            activation(),
+            nn.MaxPool2d((1, max_pool_size)),
         )
-        self.len_last_layer = self._len_last_layer(self.n_chans, self.n_times)
         self.return_feats = return_feats
 
+        dim_conv_1 = (
+            self.n_times + 2 * pad_size - (time_conv_size - 1)
+        ) // max_pool_size
+        dim_after_conv = (
+            dim_conv_1 + 2 * pad_size - (time_conv_size - 1)
+        ) // max_pool_size
+
+        self.len_last_layer = n_conv_chs * self.n_chans * dim_after_conv
+
         # TODO: Add new way to handle return_features == True
         if not return_feats:
             self.final_layer = nn.Sequential(
-                nn.Dropout(dropout),
-                nn.Linear(self.len_last_layer, self.n_outputs),
+                nn.Dropout(p=drop_prob),
+                nn.Linear(in_features=self.len_last_layer, out_features=self.n_outputs),
             )
 
-    def _len_last_layer(self, n_channels, input_size):
-        self.feature_extractor.eval()
-        with torch.no_grad():
-            out = self.feature_extractor(
-                torch.Tensor(1, 1, n_channels, input_size))
-        self.feature_extractor.train()
-        return len(out.flatten())
-
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Forward pass.
 
@@ -152,5 +153,5 @@ class SleepStagerChambon2018(EEGModuleMixin, nn.Module):
 
         if self.return_feats:
             return feats
-        else:
-            return self.final_layer(feats)
+
+        return self.final_layer(feats)