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

braindecode/models/eegnex.py

@@ -0,0 +1,247 @@
+# Authors: Bruno Aristimunha <b.aristimunha@gmail.com>
+#
+# License: BSD (3-clause)
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+from einops.layers.torch import Rearrange
+
+from braindecode.models.base import EEGModuleMixin
+from braindecode.modules import Conv2dWithConstraint, LinearWithConstraint
+
+
+class EEGNeX(EEGModuleMixin, nn.Module):
+    """EEGNeX model from Chen et al. (2024) [eegnex]_.
+
+    .. figure:: https://braindecode.org/dev/_static/model/eegnex.jpg
+        :align: center
+        :alt: EEGNeX Architecture
+
+    Parameters
+    ----------
+    activation : nn.Module, optional
+        Activation function to use. Default is `nn.ELU`.
+    depth_multiplier : int, optional
+        Depth multiplier for the depthwise convolution. Default is 2.
+    filter_1 : int, optional
+        Number of filters in the first convolutional layer. Default is 8.
+    filter_2 : int, optional
+        Number of filters in the second convolutional layer. Default is 32.
+    drop_prob: float, optional
+        Dropout rate. Default is 0.5.
+    kernel_block_4 : tuple[int, int], optional
+        Kernel size for block 4. Default is (1, 16).
+    dilation_block_4 : tuple[int, int], optional
+        Dilation rate for block 4. Default is (1, 2).
+    avg_pool_block4 : tuple[int, int], optional
+        Pooling size for block 4. Default is (1, 4).
+    kernel_block_5 : tuple[int, int], optional
+        Kernel size for block 5. Default is (1, 16).
+    dilation_block_5 : tuple[int, int], optional
+        Dilation rate for block 5. Default is (1, 4).
+    avg_pool_block5 : tuple[int, int], optional
+        Pooling size for block 5. Default is (1, 8).
+
+    Notes
+    -----
+    This implementation is not guaranteed to be correct, has not been checked
+    by original authors, only reimplemented from the paper description and
+    source code in tensorflow [EEGNexCode]_.
+
+    References
+    ----------
+    .. [eegnex] Chen, X., Teng, X., Chen, H., Pan, Y., & Geyer, P. (2024).
+       Toward reliable signals decoding for electroencephalogram: A benchmark
+       study to EEGNeX. Biomedical Signal Processing and Control, 87, 105475.
+    .. [EEGNexCode] Chen, X., Teng, X., Chen, H., Pan, Y., & Geyer, P. (2024).
+       Toward reliable signals decoding for electroencephalogram: A benchmark
+       study to EEGNeX. https://github.com/chenxiachan/EEGNeX
+    """
+
+    def __init__(
+        self,
+        # Signal related parameters
+        n_chans=None,
+        n_outputs=None,
+        n_times=None,
+        chs_info=None,
+        input_window_seconds=None,
+        sfreq=None,
+        # Model parameters
+        activation: nn.Module = nn.ELU,
+        depth_multiplier: int = 2,
+        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,
+        kernel_block_5: int = 16,
+        dilation_block_5: int = 4,
+        avg_pool_block5: int = 8,
+        max_norm_conv: float = 1.0,
+        max_norm_linear: float = 0.25,
+    ):
+        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.depth_multiplier = depth_multiplier
+        self.filter_1 = filter_1
+        self.filter_2 = filter_2
+        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.avg_pool_block4 = (1, avg_pool_block4)
+        self.kernel_block_5 = (1, kernel_block_5)
+        self.dilation_block_5 = (1, dilation_block_5)
+        self.avg_pool_block5 = (1, avg_pool_block5)
+
+        # final layers output
+        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=self.kernel_block_1_2,
+                padding="same",
+                bias=False,
+            ),
+            nn.BatchNorm2d(num_features=self.filter_1),
+        )
+
+        self.block_2 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=self.filter_1,
+                out_channels=self.filter_2,
+                kernel_size=self.kernel_block_1_2,
+                padding="same",
+                bias=False,
+            ),
+            nn.BatchNorm2d(num_features=self.filter_2),
+        )
+
+        self.block_3 = nn.Sequential(
+            Conv2dWithConstraint(
+                in_channels=self.filter_2,
+                out_channels=self.filter_3,
+                max_norm=max_norm_conv,
+                kernel_size=(self.n_chans, 1),
+                groups=self.filter_2,
+                bias=False,
+            ),
+            nn.BatchNorm2d(num_features=self.filter_3),
+            self.activation(),
+            nn.AvgPool2d(
+                kernel_size=self.avg_pool_block4,
+                padding=(0, 1),
+            ),
+            nn.Dropout(p=self.drop_prob),
+        )
+
+        self.block_4 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=self.filter_3,
+                out_channels=self.filter_2,
+                kernel_size=self.kernel_block_4,
+                dilation=self.dilation_block_4,
+                padding="same",
+                bias=False,
+            ),
+            nn.BatchNorm2d(num_features=self.filter_2),
+        )
+
+        self.block_5 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=self.filter_2,
+                out_channels=self.filter_1,
+                kernel_size=self.kernel_block_5,
+                dilation=self.dilation_block_5,
+                padding="same",
+                bias=False,
+            ),
+            nn.BatchNorm2d(num_features=self.filter_1),
+            self.activation(),
+            nn.AvgPool2d(
+                kernel_size=self.avg_pool_block5,
+                padding=(0, 1),
+            ),
+            nn.Dropout(p=self.drop_prob),
+            nn.Flatten(),
+        )
+
+        self.final_layer = LinearWithConstraint(
+            in_features=self.in_features,
+            out_features=self.n_outputs,
+            max_norm=max_norm_linear,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the EEGNeX model.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (batch_size, n_chans, n_times).
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor of shape (batch_size, n_outputs).
+        """
+        # x shape: (batch_size, n_chans, n_times)
+        x = self.block_1(x)
+        # (batch_size, n_filter, n_chans, n_times)
+        x = self.block_2(x)
+        # (batch_size, n_filter*4, n_chans, n_times)
+        x = self.block_3(x)
+        # (batch_size, 1, n_filter*8, n_times//4)
+        x = self.block_4(x)
+        # (batch_size, 1, n_filter*8, n_times//4)
+        x = self.block_5(x)
+        # (batch_size, n_filter*(n_times//32))
+        x = self.final_layer(x)
+
+        return x
+
+    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]
+
+        # Padding for the time dimension (assumed from padding=(0, 1))
+        pad4 = 1
+        pad5 = 1
+
+        # 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/eegresnet.py

@@ -0,0 +1,362 @@
+# Authors: Robin Tibor Schirrmeister <robintibor@gmail.com>
+#          Tonio Ball
+#
+# License: BSD-3
+
+import numpy as np
+import torch
+from einops.layers.torch import Rearrange
+from torch import nn
+from torch.nn import init
+
+from braindecode.models.base import EEGModuleMixin
+from braindecode.modules import (
+    AvgPool2dWithConv,
+    Ensure4d,
+    SqueezeFinalOutput,
+)
+
+
+class EEGResNet(EEGModuleMixin, nn.Sequential):
+    """EEGResNet from Schirrmeister et al. 2017 [Schirrmeister2017]_.
+
+    .. figure:: https://onlinelibrary.wiley.com/cms/asset/bed1b768-809f-4bc6-b942-b36970d81271/hbm23730-fig-0003-m.jpg
+        :align: center
+        :alt: EEGResNet Architecture
+
+    Model described in [Schirrmeister2017]_.
+
+    Parameters
+    ----------
+    in_chans :
+        Alias for ``n_chans``.
+    n_classes :
+        Alias for ``n_outputs``.
+    input_window_samples :
+       Alias for ``n_times``.
+    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``.
+
+    References
+    ----------
+    .. [Schirrmeister2017] Schirrmeister, R. T., Springenberg, J. T., Fiederer,
+       L. D. J., Glasstetter, M., Eggensperger, K., Tangermann, M., Hutter, F.
+       & Ball, T. (2017). Deep learning with convolutional neural networks for ,
+       EEG decoding and visualization. Human Brain Mapping, Aug. 2017.
+       Online: http://dx.doi.org/10.1002/hbm.23730
+    """
+
+    def __init__(
+        self,
+        n_chans=None,
+        n_outputs=None,
+        n_times=None,
+        final_pool_length="auto",
+        n_first_filters=20,
+        n_layers_per_block=2,
+        first_filter_length=3,
+        activation=nn.ELU,
+        split_first_layer=True,
+        batch_norm_alpha=0.1,
+        batch_norm_epsilon=1e-4,
+        conv_weight_init_fn=lambda w: init.kaiming_normal_(w, a=0),
+        chs_info=None,
+        input_window_seconds=None,
+        sfreq=250,
+    ):
+        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
+
+        if final_pool_length == "auto":
+            assert self.n_times is not None
+        assert first_filter_length % 2 == 1
+        self.final_pool_length = final_pool_length
+        self.n_first_filters = n_first_filters
+        self.n_layers_per_block = n_layers_per_block
+        self.first_filter_length = first_filter_length
+        self.nonlinearity = activation
+        self.split_first_layer = split_first_layer
+        self.batch_norm_alpha = batch_norm_alpha
+        self.batch_norm_epsilon = batch_norm_epsilon
+        self.conv_weight_init_fn = conv_weight_init_fn
+
+        self.mapping = {
+            "conv_classifier.weight": "final_layer.conv_classifier.weight",
+            "conv_classifier.bias": "final_layer.conv_classifier.bias",
+        }
+
+        self.add_module("ensuredims", Ensure4d())
+        if self.split_first_layer:
+            self.add_module("dimshuffle", Rearrange("batch C T 1 -> batch 1 T C"))
+            self.add_module(
+                "conv_time",
+                nn.Conv2d(
+                    1,
+                    self.n_first_filters,
+                    (self.first_filter_length, 1),
+                    stride=1,
+                    padding=(self.first_filter_length // 2, 0),
+                ),
+            )
+            self.add_module(
+                "conv_spat",
+                nn.Conv2d(
+                    self.n_first_filters,
+                    self.n_first_filters,
+                    (1, self.n_chans),
+                    stride=(1, 1),
+                    bias=False,
+                ),
+            )
+        else:
+            self.add_module(
+                "conv_time",
+                nn.Conv2d(
+                    self.n_chans,
+                    self.n_first_filters,
+                    (self.first_filter_length, 1),
+                    stride=(1, 1),
+                    padding=(self.first_filter_length // 2, 0),
+                    bias=False,
+                ),
+            )
+        n_filters_conv = self.n_first_filters
+        self.add_module(
+            "bnorm",
+            nn.BatchNorm2d(
+                n_filters_conv, momentum=self.batch_norm_alpha, affine=True, eps=1e-5
+            ),
+        )
+        self.add_module("conv_nonlin", self.nonlinearity())
+        cur_dilation = np.array([1, 1])
+        n_cur_filters = n_filters_conv
+        i_block = 1
+        for i_layer in range(self.n_layers_per_block):
+            self.add_module(
+                "res_{:d}_{:d}".format(i_block, i_layer),
+                _ResidualBlock(n_cur_filters, n_cur_filters, dilation=cur_dilation),
+            )
+        i_block += 1
+        cur_dilation[0] *= 2
+        n_out_filters = int(2 * n_cur_filters)
+        self.add_module(
+            "res_{:d}_{:d}".format(i_block, 0),
+            _ResidualBlock(
+                n_cur_filters,
+                n_out_filters,
+                dilation=cur_dilation,
+            ),
+        )
+        n_cur_filters = n_out_filters
+        for i_layer in range(1, self.n_layers_per_block):
+            self.add_module(
+                "res_{:d}_{:d}".format(i_block, i_layer),
+                _ResidualBlock(n_cur_filters, n_cur_filters, dilation=cur_dilation),
+            )
+
+        i_block += 1
+        cur_dilation[0] *= 2
+        n_out_filters = int(1.5 * n_cur_filters)
+        self.add_module(
+            "res_{:d}_{:d}".format(i_block, 0),
+            _ResidualBlock(
+                n_cur_filters,
+                n_out_filters,
+                dilation=cur_dilation,
+            ),
+        )
+        n_cur_filters = n_out_filters
+        for i_layer in range(1, self.n_layers_per_block):
+            self.add_module(
+                "res_{:d}_{:d}".format(i_block, i_layer),
+                _ResidualBlock(n_cur_filters, n_cur_filters, dilation=cur_dilation),
+            )
+
+        i_block += 1
+        cur_dilation[0] *= 2
+        self.add_module(
+            "res_{:d}_{:d}".format(i_block, 0),
+            _ResidualBlock(
+                n_cur_filters,
+                n_cur_filters,
+                dilation=cur_dilation,
+            ),
+        )
+        for i_layer in range(1, self.n_layers_per_block):
+            self.add_module(
+                "res_{:d}_{:d}".format(i_block, i_layer),
+                _ResidualBlock(n_cur_filters, n_cur_filters, dilation=cur_dilation),
+            )
+
+        i_block += 1
+        cur_dilation[0] *= 2
+        self.add_module(
+            "res_{:d}_{:d}".format(i_block, 0),
+            _ResidualBlock(
+                n_cur_filters,
+                n_cur_filters,
+                dilation=cur_dilation,
+            ),
+        )
+        for i_layer in range(1, self.n_layers_per_block):
+            self.add_module(
+                "res_{:d}_{:d}".format(i_block, i_layer),
+                _ResidualBlock(n_cur_filters, n_cur_filters, dilation=cur_dilation),
+            )
+
+        i_block += 1
+        cur_dilation[0] *= 2
+        self.add_module(
+            "res_{:d}_{:d}".format(i_block, 0),
+            _ResidualBlock(
+                n_cur_filters,
+                n_cur_filters,
+                dilation=cur_dilation,
+            ),
+        )
+        for i_layer in range(1, self.n_layers_per_block):
+            self.add_module(
+                "res_{:d}_{:d}".format(i_block, i_layer),
+                _ResidualBlock(n_cur_filters, n_cur_filters, dilation=cur_dilation),
+            )
+        i_block += 1
+        cur_dilation[0] *= 2
+        self.add_module(
+            "res_{:d}_{:d}".format(i_block, 0),
+            _ResidualBlock(
+                n_cur_filters,
+                n_cur_filters,
+                dilation=cur_dilation,
+            ),
+        )
+        for i_layer in range(1, self.n_layers_per_block):
+            self.add_module(
+                "res_{:d}_{:d}".format(i_block, i_layer),
+                _ResidualBlock(n_cur_filters, n_cur_filters, dilation=cur_dilation),
+            )
+
+        self.eval()
+        if self.final_pool_length == "auto":
+            self.add_module("mean_pool", nn.AdaptiveAvgPool2d((1, 1)))
+        else:
+            pool_dilation = int(cur_dilation[0]), int(cur_dilation[1])
+            self.add_module(
+                "mean_pool",
+                AvgPool2dWithConv(
+                    (self.final_pool_length, 1), (1, 1), dilation=pool_dilation
+                ),
+            )
+
+        # Incorporating classification module and subsequent ones in one final layer
+        module = nn.Sequential()
+
+        module.add_module(
+            "conv_classifier",
+            nn.Conv2d(
+                n_cur_filters,
+                self.n_outputs,
+                (1, 1),
+                bias=True,
+            ),
+        )
+
+        module.add_module("squeeze", SqueezeFinalOutput())
+
+        self.add_module("final_layer", module)
+
+        # Initialize all weights
+        self.apply(lambda module: self._weights_init(module, self.conv_weight_init_fn))
+
+        # Start in train mode
+        self.train()
+
+    @staticmethod
+    def _weights_init(module, conv_weight_init_fn):
+        """
+        initialize weights
+        """
+        classname = module.__class__.__name__
+        if "Conv" in classname and classname != "AvgPool2dWithConv":
+            conv_weight_init_fn(module.weight)
+            if module.bias is not None:
+                init.constant_(module.bias, 0)
+        elif "BatchNorm" in classname:
+            init.constant_(module.weight, 1)
+            init.constant_(module.bias, 0)
+
+
+class _ResidualBlock(nn.Module):
+    """
+    create a residual learning building block with two stacked 3x3 convlayers as in paper
+    """
+
+    def __init__(
+        self,
+        in_filters,
+        out_num_filters,
+        dilation,
+        filter_time_length=3,
+        nonlinearity: nn.Module = nn.ELU,
+        batch_norm_alpha=0.1,
+        batch_norm_epsilon=1e-4,
+    ):
+        super(_ResidualBlock, self).__init__()
+        time_padding = int((filter_time_length - 1) * dilation[0])
+        assert time_padding % 2 == 0
+        time_padding = int(time_padding // 2)
+        dilation = (int(dilation[0]), int(dilation[1]))
+        assert (out_num_filters - in_filters) % 2 == 0, (
+            "Need even number of extra channels in order to be able to pad correctly"
+        )
+        self.n_pad_chans = out_num_filters - in_filters
+
+        self.conv_1 = nn.Conv2d(
+            in_filters,
+            out_num_filters,
+            (filter_time_length, 1),
+            stride=(1, 1),
+            dilation=dilation,
+            padding=(time_padding, 0),
+        )
+        self.bn1 = nn.BatchNorm2d(
+            out_num_filters,
+            momentum=batch_norm_alpha,
+            affine=True,
+            eps=batch_norm_epsilon,
+        )
+        self.conv_2 = nn.Conv2d(
+            out_num_filters,
+            out_num_filters,
+            (filter_time_length, 1),
+            stride=(1, 1),
+            dilation=dilation,
+            padding=(time_padding, 0),
+        )
+        self.bn2 = nn.BatchNorm2d(
+            out_num_filters,
+            momentum=batch_norm_alpha,
+            affine=True,
+            eps=batch_norm_epsilon,
+        )
+        # also see https://mail.google.com/mail/u/0/#search/ilya+joos/1576137dd34c3127
+        # for resnet options as ilya used them
+        self.nonlinearity = nonlinearity()
+
+    def forward(self, x):
+        stack_1 = self.nonlinearity(self.bn1(self.conv_1(x)))
+        stack_2 = self.bn2(self.conv_2(stack_1))  # next nonlin after sum
+        if self.n_pad_chans != 0:
+            zeros_for_padding = x.new_zeros(
+                (x.shape[0], self.n_pad_chans // 2, x.shape[2], x.shape[3])
+            )
+            x = torch.cat((zeros_for_padding, x, zeros_for_padding), dim=1)
+        out = self.nonlinearity(x + stack_2)
+        return out