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

braindecode/models/eegsimpleconv.py

@@ -0,0 +1,199 @@
+"""
+EEG-SimpleConv is a 1D Convolutional Neural Network from Yassine El Ouahidi et al. (2023).
+
+Originally designed for Motor Imagery decoding, from EEG signals.
+The model offers competitive performances, with a low latency and is mainly composed of
+1D convolutional layers.
+
+"""
+
+# Authors: Yassine El Ouahidi <eloua.yas@gmail.com>
+#
+# License: BSD-3
+
+import torch
+from torch import nn
+from torchaudio.transforms import Resample
+
+from braindecode.models.base import EEGModuleMixin
+
+
+class EEGSimpleConv(EEGModuleMixin, torch.nn.Module):
+    """EEGSimpleConv from Ouahidi, YE et al. (2023) [Yassine2023]_.
+
+    .. figure:: https://raw.githubusercontent.com/elouayas/EEGSimpleConv/refs/heads/main/architecture.png
+        :align: center
+        :alt: EEGSimpleConv Architecture
+
+    EEGSimpleConv is a 1D Convolutional Neural Network originally designed
+    for decoding motor imagery from EEG signals. The model aims to have a
+    very simple and straightforward architecture that allows a low latency,
+    while still achieving very competitive performance.
+
+    EEG-SimpleConv starts with a 1D convolutional layer, where each EEG channel
+    enters a separate 1D convolutional channel. This is followed by a series of
+    blocks of two 1D convolutional layers. Between the two convolutional layers
+    of each block is a max pooling layer, which downsamples the data by a factor
+    of 2. Each convolution is followed by a batch normalisation layer and a ReLU
+    activation function. Finally, a global average pooling (in the time domain)
+    is performed to obtain a single value per feature map, which is then fed
+    into a linear layer to obtain the final classification prediction output.
+
+    The paper and original code with more details about the methodological
+    choices are available at the [Yassine2023]_ and [Yassine2023Code]_.
+
+    The input shape should be three-dimensional matrix representing the EEG
+    signals.
+
+    ``(batch_size, n_channels, n_timesteps)``.
+
+    Notes
+    -----
+    The authors recommend using the default parameters for MI decoding.
+    Please refer to the original paper and code for more details.
+
+    Recommended range for the choice of the hyperparameters, regarding the
+    evaluation paradigm.
+
+    |    Parameter    | Within-Subject | Cross-Subject |
+    | feature_maps    | [64-144]       |   [64-144]    |
+    | n_convs         |    1           |   [2-4]       |
+    | resampling_freq | [70-100]       |   [50-80]     |
+    | kernel_size     | [12-17]        |   [5-8]       |
+
+
+    An intensive ablation study is included in the paper to understand the
+    of each parameter on the model performance.
+
+    .. versionadded:: 0.9
+
+    Parameters
+    ----------
+    feature_maps: int
+        Number of Feature Maps at the first Convolution, width of the model.
+    n_convs: int
+        Number of blocks of convolutions (2 convolutions per block), depth of the model.
+    resampling: int
+        Resampling Frequency.
+    kernel_size: int
+        Size of the convolutions kernels.
+    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
+    ----------
+    .. [Yassine2023] Yassine El Ouahidi, V. Gripon, B. Pasdeloup, G. Bouallegue
+        N. Farrugia, G. Lioi, 2023. A Strong and Simple Deep Learning Baseline for
+        BCI Motor Imagery Decoding. Arxiv preprint. arxiv.org/abs/2309.07159
+    .. [Yassine2023Code] Yassine El Ouahidi, V. Gripon, B. Pasdeloup, G. Bouallegue
+        N. Farrugia, G. Lioi, 2023. A Strong and Simple Deep Learning Baseline for
+        BCI Motor Imagery Decoding. GitHub repository.
+        https://github.com/elouayas/EEGSimpleConv.
+
+    """
+
+    def __init__(
+        self,
+        # Base arguments
+        n_outputs=None,
+        n_chans=None,
+        sfreq=None,
+        # Model specific arguments
+        feature_maps=128,
+        n_convs=2,
+        resampling_freq=80,
+        kernel_size=8,
+        return_feature=False,
+        activation: nn.Module = nn.ReLU,
+        # Other ways to initialize the model
+        chs_info=None,
+        n_times=None,
+        input_window_seconds=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, sfreq, input_window_seconds
+
+        self.return_feature = return_feature
+        self.resample = (
+            Resample(orig_freq=int(self.sfreq), new_freq=int(resampling_freq))
+            if self.sfreq != resampling_freq
+            else torch.nn.Identity()
+        )
+
+        self.conv = torch.nn.Conv1d(
+            self.n_chans,
+            feature_maps,
+            kernel_size=kernel_size,
+            padding=kernel_size // 2,
+            bias=False,
+        )
+        self.bn = torch.nn.BatchNorm1d(feature_maps)
+        self.blocks = []
+        new_feature_maps = feature_maps
+        old_feature_maps = feature_maps
+        for i in range(n_convs):
+            if i > 0:
+                # 1.414 = sqrt(2) allow constant flops.
+                new_feature_maps = int(1.414 * new_feature_maps)
+            self.blocks.append(
+                torch.nn.Sequential(
+                    (
+                        torch.nn.Conv1d(
+                            old_feature_maps,
+                            new_feature_maps,
+                            kernel_size=kernel_size,
+                            padding=kernel_size // 2,
+                            bias=False,
+                        )
+                    ),
+                    (torch.nn.BatchNorm1d(new_feature_maps)),
+                    (torch.nn.MaxPool1d(2) if i > 0 - 1 else torch.nn.MaxPool1d(1)),
+                    (activation()),
+                    (
+                        torch.nn.Conv1d(
+                            new_feature_maps,
+                            new_feature_maps,
+                            kernel_size=kernel_size,
+                            padding=kernel_size // 2,
+                            bias=False,
+                        )
+                    ),
+                    (torch.nn.BatchNorm1d(new_feature_maps)),
+                    (activation()),
+                )
+            )
+            old_feature_maps = new_feature_maps
+        self.blocks = torch.nn.ModuleList(self.blocks)
+        self.final_layer = torch.nn.Linear(old_feature_maps, self.n_outputs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the model.
+
+        Parameters
+        ----------
+        x: PyTorch Tensor
+            Input tensor of shape (batch_size, n_channels, n_times)
+
+        Returns
+        -------
+        PyTorch Tensor (optional)
+            Output tensor of shape (batch_size, n_outputs)
+        """
+        x_rs = self.resample(x.contiguous())
+        feat = torch.relu(self.bn(self.conv(x_rs)))
+        for seq in self.blocks:
+            feat = seq(feat)
+        feat = feat.mean(dim=2)
+        if self.return_feature:
+            return feat
+        else:
+            return self.final_layer(feat)

braindecode/models/eegtcnet.py

@@ -0,0 +1,335 @@
+# Authors: Bruno Aristimunha <b.aristimunha@gmail.com>
+#
+# License: BSD (3-clause)
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from einops.layers.torch import Rearrange
+
+from braindecode.models.base import EEGModuleMixin
+from braindecode.modules import Chomp1d, MaxNormLinear
+
+
+class EEGTCNet(EEGModuleMixin, nn.Module):
+    """EEGTCNet model from Ingolfsson et al. (2020) [ingolfsson2020]_.
+
+    .. figure:: https://braindecode.org/dev/_static/model/eegtcnet.jpg
+        :align: center
+        :alt: EEGTCNet Architecture
+
+    Combining EEGNet and TCN blocks.
+
+    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 temporal filters in the first convolutional layer. Default is 8.
+    kern_length : int, optional
+        Length of the temporal kernel in the first convolutional layer. Default is 64.
+    dropout : float, optional
+        Dropout rate. Default is 0.5.
+    depth : int, optional
+        Number of residual blocks in the TCN. Default is 2.
+    kernel_size : int, optional
+        Size of the temporal convolutional kernel in the TCN. Default is 4.
+    filters : int, optional
+        Number of filters in the TCN convolutional layers. Default is 12.
+    max_norm_const : float
+        Maximum L2-norm constraint imposed on weights of the last
+        fully-connected layer. Defaults to 0.25.
+
+    References
+    ----------
+    .. [ingolfsson2020] Ingolfsson, T. M., Hersche, M., Wang, X., Kobayashi, N.,
+        Cavigelli, L., & Benini, L. (2020). EEG-TCNet: An accurate temporal
+        convolutional network for embedded motor-imagery brain–machine interfaces.
+        https://doi.org/10.48550/arXiv.2006.00622
+    """
+
+    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,
+        kern_length: int = 64,
+        drop_prob: float = 0.5,
+        depth: int = 2,
+        kernel_size: int = 4,
+        filters: int = 12,
+        max_norm_const: 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.activation = activation
+        self.drop_prob = drop_prob
+        self.depth_multiplier = depth_multiplier
+        self.filter_1 = filter_1
+        self.kern_length = kern_length
+        self.depth = depth
+        self.kernel_size = kernel_size
+        self.filters = filters
+        self.max_norm_const = max_norm_const
+        self.filter_2 = self.filter_1 * self.depth_multiplier
+
+        self.arrange_dim_input = Rearrange(
+            "batch nchans ntimes -> batch 1 ntimes nchans"
+        )
+        # EEGNet_TC Block
+        self.eegnet_tc = _EEGNetTC(
+            n_chans=self.n_chans,
+            filter_1=self.filter_1,
+            kern_length=self.kern_length,
+            depth_multiplier=self.depth_multiplier,
+            drop_prob=self.drop_prob,
+            activation=self.activation,
+        )
+        self.arrange_dim_eegnet = Rearrange(
+            "batch filter2 rtimes 1 -> batch rtimes filter2"
+        )
+
+        # TCN Block
+        self.tcn_block = _TCNBlock(
+            input_dimension=self.filter_2,
+            depth=self.depth,
+            kernel_size=self.kernel_size,
+            filters=self.filters,
+            drop_prob=self.drop_prob,
+            activation=self.activation,
+        )
+
+        # Classification Block
+        self.final_layer = MaxNormLinear(
+            in_features=self.filters,
+            out_features=self.n_outputs,
+            max_norm_val=self.max_norm_const,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the EEGTCNet 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.arrange_dim_input(x)  # (batch_size, 1, n_times, n_chans)
+        x = self.eegnet_tc(x)  # (batch_size, filter, reduced_time, 1)
+
+        x = self.arrange_dim_eegnet(x)  # (batch_size, reduced_time, F2)
+        x = self.tcn_block(x)  # (batch_size, time_steps, filters)
+
+        # Select the last time step
+        x = x[:, -1, :]  # (batch_size, filters)
+
+        x = self.final_layer(x)  # (batch_size, n_outputs)
+
+        return x
+
+
+class _EEGNetTC(nn.Module):
+    """EEGNet Temporal Convolutional Network (TCN) block.
+
+    The main difference from our EEGNetV4 (braindecode) implementation is the
+    kernel and dimensional order. Because of this, we decided to keep this
+    implementation in a future issue; we will re-evaluate if it is necessary
+    to maintain this separate implementation.
+
+    Parameters
+    ----------
+    n_chans : int
+        Number of EEG channels.
+    filter_1 : int
+        Number of temporal filters in the first convolutional layer.
+    kern_length : int
+        Length of the temporal kernel in the first convolutional layer.
+    depth_multiplier : int
+        Depth multiplier for the depthwise convolution.
+    drop_prob : float
+        Dropout rate.
+    activation : nn.Module
+        Activation function.
+    """
+
+    def __init__(
+        self,
+        n_chans: int,
+        filter_1: int = 8,
+        kern_length: int = 64,
+        depth_multiplier: int = 2,
+        drop_prob: float = 0.5,
+        activation: nn.Module = nn.ELU,
+    ):
+        super().__init__()
+        self.activation = activation()
+        self.drop_prob = drop_prob
+        self.n_chans = n_chans
+        self.filter_1 = filter_1
+        self.filter_2 = self.filter_1 * depth_multiplier
+
+        # First Conv2D Layer
+        self.conv1 = nn.Conv2d(
+            in_channels=1,
+            out_channels=self.filter_1,
+            kernel_size=(kern_length, 1),
+            padding=(kern_length // 2, 0),
+            bias=False,
+        )
+        self.bn1 = nn.BatchNorm2d(self.filter_1)
+
+        # Depthwise Convolution
+        self.depthwise_conv = nn.Conv2d(
+            in_channels=self.filter_1,
+            out_channels=self.filter_2,
+            kernel_size=(1, n_chans),
+            groups=self.filter_1,
+            bias=False,
+        )
+        self.bn2 = nn.BatchNorm2d(self.filter_2)
+        self.pool1 = nn.AvgPool2d(kernel_size=(8, 1))
+        self.drop1 = nn.Dropout(p=drop_prob)
+
+        # Separable Convolution (Depthwise + Pointwise)
+        self.separable_conv_depthwise = nn.Conv2d(
+            in_channels=self.filter_2,
+            out_channels=self.filter_2,
+            kernel_size=(self.filter_2, 1),
+            groups=self.filter_2,
+            padding=(self.filter_2 // 2, 0),
+            bias=False,
+        )
+        self.separable_conv_pointwise = nn.Conv2d(
+            in_channels=self.filter_2,
+            out_channels=self.filter_2,
+            kernel_size=(1, 1),
+            bias=False,
+        )
+        self.bn3 = nn.BatchNorm2d(self.filter_2)
+        self.pool2 = nn.AvgPool2d(kernel_size=(self.filter_1, 1))
+        self.drop2 = nn.Dropout(p=drop_prob)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x shape: (batch_size, 1, n_times, n_chans)
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.activation(x)
+
+        x = self.depthwise_conv(x)
+        x = self.bn2(x)
+        x = self.activation(x)
+        x = self.pool1(x)
+        x = self.drop1(x)
+
+        x = self.separable_conv_depthwise(x)
+        x = self.separable_conv_pointwise(x)
+        x = self.bn3(x)
+        x = self.activation(x)
+        x = self.pool2(x)
+        x = self.drop2(x)
+
+        return x  # Shape: (batch_size, F2, reduced_time, 1)
+
+
+class _TCNBlock(nn.Module):
+    """
+    Many differences from our Temporal Block (braindecode) implementation.
+    Because of this, we decided to keep this implementation in a future issue;
+    we will re-evaluate if it is necessary to maintain this separate
+    implementation.
+
+
+    """
+
+    def __init__(
+        self,
+        input_dimension: int,
+        depth: int,
+        kernel_size: int,
+        filters: int,
+        drop_prob: float,
+        activation: nn.Module = nn.ELU,
+    ):
+        super().__init__()
+        self.activation = activation()
+        self.drop_prob = drop_prob
+        self.depth = depth
+        self.filters = filters
+        self.kernel_size = kernel_size
+
+        self.layers = nn.ModuleList()
+        self.downsample = (
+            nn.Conv1d(input_dimension, filters, kernel_size=1, bias=False)
+            if input_dimension != filters
+            else None
+        )
+
+        for i in range(depth):
+            dilation = 2**i
+            padding = (kernel_size - 1) * dilation
+            conv_block = nn.Sequential(
+                nn.Conv1d(
+                    in_channels=input_dimension if i == 0 else filters,
+                    out_channels=filters,
+                    kernel_size=kernel_size,
+                    dilation=dilation,
+                    padding=padding,
+                    bias=False,
+                ),
+                Chomp1d(padding),
+                self.activation,
+                nn.Dropout(self.drop_prob),
+                nn.Conv1d(
+                    in_channels=filters,
+                    out_channels=filters,
+                    kernel_size=kernel_size,
+                    dilation=dilation,
+                    padding=padding,
+                    bias=False,
+                ),
+                Chomp1d(padding),
+                self.activation,
+                nn.Dropout(self.drop_prob),
+            )
+            self.layers.append(conv_block)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x shape: (batch_size, time_steps, input_dimension)
+        x = x.permute(0, 2, 1)  # (batch_size, input_dimension, time_steps)
+
+        res = x if self.downsample is None else self.downsample(x)
+        for layer in self.layers:
+            out = layer(x)
+            out = out + res
+            out = self.activation(out)
+            res = out  # Update residual
+            x = out  # Update input for next layer
+
+        out = out.permute(0, 2, 1)  # (batch_size, time_steps, filters)
+        return out