braindecode 1.3.0.dev177069446__py3-none-any.whl

braindecode/models/eegtcnet.py

@@ -0,0 +1,337 @@
+# 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):
+    r"""EEGTCNet model from Ingolfsson et al (2020) [ingolfsson2020]_.
+
+    :bdg-success:`Convolution` :bdg-secondary:`Recurrent`
+
+    .. 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: type[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):
+    r"""EEGNet Temporal Convolutional Network (TCN) block.
+
+    The main difference from our :class:`EEGNet` (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: type[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):
+    r"""
+    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: type[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

braindecode/models/fbcnet.py

@@ -0,0 +1,225 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from einops.layers.torch import Rearrange
+from mne.utils import warn
+from torch import nn
+
+from braindecode.models.base import EEGModuleMixin
+from braindecode.modules import (
+    Conv2dWithConstraint,
+    FilterBankLayer,
+    LinearWithConstraint,
+    LogVarLayer,
+    MaxLayer,
+    MeanLayer,
+    StdLayer,
+    VarLayer,
+)
+
+_valid_layers = {
+    "VarLayer": VarLayer,
+    "StdLayer": StdLayer,
+    "LogVarLayer": LogVarLayer,
+    "MeanLayer": MeanLayer,
+    "MaxLayer": MaxLayer,
+}
+
+
+class FBCNet(EEGModuleMixin, nn.Module):
+    r"""FBCNet from Mane, R et al (2021) [fbcnet2021]_.
+
+    :bdg-success:`Convolution` :bdg-primary:`Filterbank`
+
+    .. figure:: https://raw.githubusercontent.com/ravikiran-mane/FBCNet/refs/heads/master/FBCNet-V2.png
+        :align: center
+        :alt: FBCNet Architecture
+
+    The FBCNet model applies spatial convolution and variance calculation along
+    the time axis, inspired by the Filter Bank Common Spatial Pattern (FBCSP)
+    algorithm.
+
+    Notes
+    -----
+    This implementation is not guaranteed to be correct and has not been checked
+    by the original authors; it has only been reimplemented from the paper
+    description and source code [fbcnetcode2021]_. There is a difference in the
+    activation function; in the paper, the ELU is used as the activation function,
+    but in the original code, SiLU is used. We followed the code.
+
+    Parameters
+    ----------
+    n_bands : int or None or list[tuple[int, int]]], default=9
+        Number of frequency bands. Could
+    n_filters_spat : int, default=32
+        Number of spatial filters for the first convolution.
+    n_dim: int, default=3
+        Number of dimensions for the temporal reductor
+    temporal_layer : str, default='LogVarLayer'
+        Type of temporal aggregator layer. Options: 'VarLayer', 'StdLayer',
+        'LogVarLayer', 'MeanLayer', 'MaxLayer'.
+    stride_factor : int, default=4
+        Stride factor for reshaping.
+    activation : nn.Module, default=nn.SiLU
+        Activation function class to apply in Spatial Convolution Block.
+    cnn_max_norm : float, default=2.0
+        Maximum norm for the spatial convolution layer.
+    linear_max_norm : float, default=0.5
+        Maximum norm for the final linear layer.
+    filter_parameters: dict, default None
+        Dictionary of parameters to use for the FilterBankLayer.
+        If None, a default Chebyshev Type II filter with transition bandwidth of
+        2 Hz and stop-band ripple of 30 dB will be used.
+
+    References
+    ----------
+    .. [fbcnet2021] Mane, R., Chew, E., Chua, K., Ang, K. K., Robinson, N.,
+        Vinod, A. P., ... & Guan, C. (2021). FBCNet: A multi-view convolutional
+        neural network for brain-computer interface. preprint arXiv:2104.01233.
+    .. [fbcnetcode2021] Link to source-code:
+        https://github.com/ravikiran-mane/FBCNet
+    """
+
+    def __init__(
+        self,
+        # Braindecode parameters
+        n_chans=None,
+        n_outputs=None,
+        chs_info=None,
+        n_times=None,
+        input_window_seconds=None,
+        sfreq=None,
+        # models parameters
+        n_bands=9,
+        n_filters_spat: int = 32,
+        temporal_layer: str = "LogVarLayer",
+        n_dim: int = 3,
+        stride_factor: int = 4,
+        activation: type[nn.Module] = nn.SiLU,
+        linear_max_norm: float = 0.5,
+        cnn_max_norm: float = 2.0,
+        filter_parameters: dict[Any, Any] | None = None,
+    ):
+        super().__init__(
+            n_chans=n_chans,
+            n_outputs=n_outputs,
+            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
+
+        # Parameters
+        self.n_bands = n_bands
+        self.n_filters_spat = n_filters_spat
+        self.n_dim = n_dim
+        self.stride_factor = stride_factor
+        self.activation = activation
+        self.filter_parameters = filter_parameters or {}
+
+        # Checkers
+        if temporal_layer not in _valid_layers:
+            raise NotImplementedError(
+                f"Temporal layer '{temporal_layer}' is not implemented."
+            )
+
+        if self.n_times % self.stride_factor != 0:
+            warn(
+                f"Time dimension ({self.n_times}) is not divisible by"
+                f" stride_factor ({self.stride_factor}). Input will be padded.",
+                UserWarning,
+            )
+
+        # Layers
+        # Following paper nomenclature
+        self.spectral_filtering = FilterBankLayer(
+            n_chans=self.n_chans,
+            sfreq=self.sfreq,
+            band_filters=self.n_bands,
+            verbose=False,
+            **self.filter_parameters,
+        )
+        # As we have an internal process to create the bands,
+        # we get the values from the filterbank
+        self.n_bands = self.spectral_filtering.n_bands
+
+        # Spatial Convolution Block (SCB)
+        self.spatial_conv = nn.Sequential(
+            Conv2dWithConstraint(
+                in_channels=self.n_bands,
+                out_channels=self.n_filters_spat * self.n_bands,
+                kernel_size=(self.n_chans, 1),
+                groups=self.n_bands,
+                max_norm=cnn_max_norm,
+                padding=0,
+            ),
+            nn.BatchNorm2d(self.n_filters_spat * self.n_bands),
+            self.activation(),
+        )
+
+        # Padding layer
+        if self.n_times % self.stride_factor != 0:
+            self.padding_size = stride_factor - (self.n_times % stride_factor)
+            self.n_times_padded = self.n_times + self.padding_size
+            self.padding_layer = nn.ConstantPad1d((0, self.padding_size), 0.0)
+        else:
+            self.padding_layer = nn.Identity()
+            self.n_times_padded = self.n_times
+
+        # Temporal aggregator
+        self.temporal_layer = _valid_layers[temporal_layer](dim=self.n_dim)  # type: ignore
+
+        # Flatten layer
+        self.flatten_layer = Rearrange("batch ... -> batch (...)")
+
+        # Final fully connected layer
+        self.final_layer = LinearWithConstraint(
+            in_features=self.n_filters_spat * self.n_bands * self.stride_factor,
+            out_features=self.n_outputs,
+            max_norm=linear_max_norm,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the FBCNet model.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor with shape (batch_size, n_chans, n_times).
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor with shape (batch_size, n_outputs).
+        """
+        # output: (batch_size, n_chans, n_times)
+        x = self.spectral_filtering(x)
+
+        # output: (batch_size, n_bands, n_chans, n_times)
+        x = self.spatial_conv(x)
+        batch_size, channels, _, _ = x.shape
+
+        # shape: (batch_size, n_filters_spat * n_bands, 1, n_times)
+        x = self.padding_layer(x)
+
+        # shape: (batch_size, n_filters_spat * n_bands, 1, n_times_padded)
+        x = x.view(
+            batch_size,
+            channels,
+            self.stride_factor,
+            self.n_times_padded // self.stride_factor,
+        )
+        # shape: batch_size, n_filters_spat * n_bands, stride, n_times_padded/stride
+        x = self.temporal_layer(x)  # type: ignore[operator]
+
+        # shape: batch_size, n_filters_spat * n_bands, stride, 1
+        x = self.flatten_layer(x)
+
+        # shape: batch_size, n_filters_spat * n_bands * stride
+        x = self.final_layer(x)
+        # shape: batch_size, n_outputs
+        return x