Update SCCNet and EEGNex and enhance documentation (#742)

* updating the sccnet

* updating the whats new

* remove the commment

* using same

* expose more parameters

* fixing typo

* final fix for eeg net

* done eegnex

Author: bruAristimunha

braindecode/models/eegnex.py

@@ -3,6 +3,8 @@
 # License: BSD (3-clause)
 from __future__ import annotations
 
+import math
+
 import torch
 import torch.nn as nn
 from einops.layers.torch import Rearrange
@@ -74,6 +76,7 @@ def __init__(
         filter_1: int = 8,
         filter_2: int = 32,
         drop_prob: float = 0.5,
+        kernel_block_1_2: int = 64,
         kernel_block_4: int = 16,
         dilation_block_4: int = 2,
         avg_pool_block4: int = 4,
@@ -99,32 +102,24 @@ def __init__(
         self.filter_3 = self.filter_2 * self.depth_multiplier
         self.drop_prob = drop_prob
         self.activation = activation
-
+        self.kernel_block_1_2 = (1, kernel_block_1_2)
         self.kernel_block_4 = (1, kernel_block_4)
         self.dilation_block_4 = (1, dilation_block_4)
-        self.padding_block_4 = self._calc_padding(
-            self.kernel_block_4, self.dilation_block_4
-        )
         self.avg_pool_block4 = (1, avg_pool_block4)
-
         self.kernel_block_5 = (1, kernel_block_5)
         self.dilation_block_5 = (1, dilation_block_5)
-
-        self.padding_block_5 = self._calc_padding(
-            self.kernel_block_5, self.dilation_block_5
-        )
         self.avg_pool_block5 = (1, avg_pool_block5)
 
         # final layers output
-        self.in_features = self.filter_1 * (self.n_times // self.filter_2)
+        self.in_features = self._calculate_output_length()
 
         # Following paper nomenclature
         self.block_1 = nn.Sequential(
             Rearrange("batch ch time -> batch 1 ch time"),
             nn.Conv2d(
                 in_channels=1,
                 out_channels=self.filter_1,
-                kernel_size=(1, 64),
+                kernel_size=self.kernel_block_1_2,
                 padding="same",
                 bias=False,
             ),
@@ -135,7 +130,7 @@ def __init__(
             nn.Conv2d(
                 in_channels=self.filter_1,
                 out_channels=self.filter_2,
-                kernel_size=(1, 64),
+                kernel_size=self.kernel_block_1_2,
                 padding="same",
                 bias=False,
             ),
@@ -166,7 +161,7 @@ def __init__(
                 out_channels=self.filter_2,
                 kernel_size=self.kernel_block_4,
                 dilation=self.dilation_block_4,
-                padding=self.padding_block_4,
+                padding="same",
                 bias=False,
             ),
             nn.BatchNorm2d(num_features=self.filter_2),
@@ -178,7 +173,7 @@ def __init__(
                 out_channels=self.filter_1,
                 kernel_size=self.kernel_block_5,
                 dilation=self.dilation_block_5,
-                padding=self.padding_block_5,
+                padding="same",
                 bias=False,
             ),
             nn.BatchNorm2d(num_features=self.filter_1),
@@ -226,26 +221,27 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         return x
 
-    @staticmethod
-    def _calc_padding(
-        kernel_size: tuple[int, int], dilation: tuple[int, int]
-    ) -> tuple[int, int]:
-        """
-        Calculate padding size for 'same' convolution with dilation.
+    def _calculate_output_length(self) -> int:
+        # Pooling kernel sizes for the time dimension
+        p4 = self.avg_pool_block4[1]
+        p5 = self.avg_pool_block5[1]
 
-        Parameters
-        ----------
-        kernel_size : tuple
-            tuple containing the kernel size (height, width).
-        dilation : tuple
-            tuple containing the dilation rate (height, width).
+        # Padding for the time dimension (assumed from padding=(0, 1))
+        pad4 = 1
+        pad5 = 1
 
-        Returns
-        -------
-        tuple
-            Padding sizes (padding_height, padding_width).
-        """
-        # Calculate padding
-        padding_height = ((kernel_size[0] - 1) * dilation[0]) // 2
-        padding_width = ((kernel_size[1] - 1) * dilation[1]) // 2
-        return padding_height, padding_width
+        # Stride is assumed to be equal to kernel size (p4 and p5)
+
+        # Calculate time dimension after block 3 pooling
+        # Formula: floor((L_in + 2*padding - kernel_size) / stride) + 1
+        T3 = math.floor((self.n_times + 2 * pad4 - p4) / p4) + 1
+
+        # Calculate time dimension after block 5 pooling
+        T5 = math.floor((T3 + 2 * pad5 - p5) / p5) + 1
+
+        # Calculate final flattened features (channels * 1 * time_dim)
+        # The spatial dimension is reduced to 1 after block 3's depthwise conv.
+        final_in_features = (
+            self.filter_1 * T5
+        )  # filter_1 is the number of channels before flatten
+        return final_in_features

braindecode/models/sccnet.py

@@ -82,6 +82,7 @@ def __init__(
         n_spatial_filters_smooth: int = 20,
         drop_prob: float = 0.5,
         activation: nn.Module = LogActivation,
+        batch_norm_momentum: float = 0.1,
     ):
         super().__init__(
             n_outputs=n_outputs,
@@ -93,23 +94,15 @@ def __init__(
         )
         del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
         # Parameters
-        self.n_filters_spat = n_spatial_filters
+        self.n_spatial_filters = n_spatial_filters
         self.n_spatial_filters_smooth = n_spatial_filters_smooth
         self.drop_prob = drop_prob
 
         self.samples_100ms = int(math.floor(self.sfreq * 0.1))
         self.kernel_size_pool = int(self.sfreq * 0.5)
         # Equivalent to 0.5 seconds
 
-        # Compute the number of features for the final linear layer
-        w_out_conv2 = (
-            self.n_times - self.samples_100ms + 1  # After second conv layer
-        )
-        w_out_pool = (
-            (w_out_conv2 - self.kernel_size_pool) // self.samples_100ms + 1
-            # After pooling layer
-        )
-        num_features = self.n_spatial_filters_smooth * w_out_pool
+        num_features = self._calc_num_features()
 
         # Layers
         self.ensure_dim = Rearrange("batch nchan times -> batch 1 nchan times")
@@ -118,23 +111,27 @@ def __init__(
 
         self.spatial_conv = nn.Conv2d(
             in_channels=1,
-            out_channels=self.n_filters_spat,
+            out_channels=self.n_spatial_filters,
             kernel_size=(self.n_chans, 1),
         )
 
+        self.spatial_batch_norm = nn.BatchNorm2d(
+            self.n_spatial_filters, momentum=batch_norm_momentum
+        )
+
         self.permute = Rearrange(
             "batch filspat nchans time -> batch nchans filspat time"
         )
 
         self.spatial_filt_conv = nn.Conv2d(
             in_channels=1,
             out_channels=self.n_spatial_filters_smooth,
-            kernel_size=(self.n_filters_spat, self.samples_100ms),
-            padding=0,
+            kernel_size=(self.n_spatial_filters, self.samples_100ms),
             bias=False,
         )
-        # Momentum following keras
-        self.batch_norm = nn.BatchNorm2d(self.n_spatial_filters_smooth, momentum=0.9)
+        self.batch_norm = nn.BatchNorm2d(
+            self.n_spatial_filters_smooth, momentum=batch_norm_momentum
+        )
 
         self.dropout = nn.Dropout(self.drop_prob)
         self.temporal_smoothing = nn.AvgPool2d(
@@ -150,6 +147,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         # Shape: (batch_size, 1, n_chans, n_times)
         x = self.spatial_conv(x)
         # Shape: (batch_size, n_filters, 1, n_times)
+        x = self.spatial_batch_norm(x)
+        # Shape: (batch_size, n_filters, 1, n_times)
         x = self.permute(x)
         # Shape: (batch_size, 1, n_filters, n_times)
         x = self.spatial_filt_conv(x)
@@ -169,3 +168,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.final_layer(x)
         # Shape: (batch_size, n_outputs)
         return x
+
+    def _calc_num_features(self) -> int:
+        # Compute the number of features for the final linear layer
+        w_out_conv2 = (
+            self.n_times - self.samples_100ms + 1  # After second conv layer
+        )
+        w_out_pool = (
+            (w_out_conv2 - self.kernel_size_pool) // self.samples_100ms + 1
+            # After pooling layer
+        )
+        num_features = self.n_spatial_filters_smooth * w_out_pool
+        return num_features

docs/whats_new.rst

@@ -80,6 +80,7 @@ Enhancements
 
 Bugs
 ~~~~
+- Making the :class:`braindecode.models.SCCNet` more compatible with paper instead of source code (:gh:`742` by `Bruno Aristimunha`_)
 - Making the :class:`braindecode.models.EEGNeX` and :class:`braindecode.models.CTNet` more compatible with paper instead of source code (:gh:`740` by `Bruno Aristimunha`_)
 - Exposing extra variable to avoid problem with the parallel process (:gh:`736` by `Pierre Guetschel`_)
 - Fixing the IFNet (:gh:`739` by `Bruno Aristimunha`_)

test/unit_tests/models/test_models.py

@@ -32,6 +32,7 @@
     EEGNetv1,
     EEGNetv4,
     EEGResNet,
+    EEGNeX,
     EEGSimpleConv,
     EEGTCNet,
     FBCNet,
@@ -1445,3 +1446,31 @@ def test_initialize_weights_conv():
     assert conv.weight.std().item() <= 0.02  # Checking trunc_normal_ std
     if conv.bias is not None:
         assert torch.allclose(conv.bias, torch.zeros_like(conv.bias))
+
+
+test_cases = [
+    pytest.param(64, id="n_times=64_perfect_multiple"),
+    pytest.param(437, id="n_times=437_trace_example"), # Expect 104
+    pytest.param(95, id="n_times=95_edge_case_1"), # Expect 24
+    pytest.param(67, id="n_times=67_edge_case_2"), # Expect 16
+    pytest.param(94, id="n_times=94_edge_case_3"), # Expect 24
+]
+
+@pytest.mark.parametrize("n_times_input", test_cases)
+def test_eegnex_final_layer_in_features(n_times_input):
+    """
+    Tests if the EEGNeX model correctly calculates the 'in_features'
+    for its final linear layer during initialization, especially for
+    n_times values that are not perfect multiples of pooling factors,
+    considering the specified padding.
+    """
+    n_chans_test = 2
+    n_outputs_test = 5
+
+    model = EEGNeX(
+        n_chans=n_chans_test,
+        n_outputs=n_outputs_test,
+        n_times=n_times_input
+    )
+
+    print(model)