braindecode 1.3.0.dev177069446__py3-none-any.whl

braindecode/models/eeginception_mi.py

@@ -0,0 +1,379 @@
+# Authors: Cedric Rommel <cedric.rommel@inria.fr>
+#
+# License: BSD (3-clause)
+
+import torch
+from einops.layers.torch import Rearrange
+from torch import nn
+
+from braindecode.models.base import EEGModuleMixin
+from braindecode.modules import Ensure4d
+
+
+class EEGInceptionMI(EEGModuleMixin, nn.Module):
+    r"""EEG Inception for Motor Imagery, as proposed in Zhang et al. (2021) [1]_
+
+    :bdg-success:`Convolution`
+
+    .. figure:: https://content.cld.iop.org/journals/1741-2552/18/4/046014/revision3/jneabed81f1_hr.jpg
+        :align: center
+        :alt: EEGInceptionMI Architecture
+
+
+    The model is strongly based on the original InceptionNet for computer
+    vision. The main goal is to extract features in parallel with different
+    scales. The network has two blocks made of 3 inception modules with a skip
+    connection.
+
+    The model is fully described in [1]_.
+
+    Notes
+    -----
+    This implementation is not guaranteed to be correct, has not been checked
+    by original authors, only reimplemented bosed on the paper [1]_.
+
+    Parameters
+    ----------
+    input_window_seconds : float, optional
+        Size of the input, in seconds. Set to 4.5 s as in [1]_ for dataset
+        BCI IV 2a.
+    sfreq : float, optional
+        EEG sampling frequency in Hz. Defaults to 250 Hz as in [1]_ for dataset
+        BCI IV 2a.
+    n_convs : int, optional
+        Number of convolution per inception wide branching. Defaults to 5 as
+        in [1]_ for dataset BCI IV 2a.
+    n_filters : int, optional
+        Number of convolutional filters for all layers of this type. Set to 48
+        as in [1]_ for dataset BCI IV 2a.
+    kernel_unit_s : float, optional
+        Size in seconds of the basic 1D convolutional kernel used in inception
+        modules. Each convolutional layer in such modules have kernels of
+        increasing size, odd multiples of this value (e.g. 0.1, 0.3, 0.5, 0.7,
+        0.9 here for ``n_convs=5``). Defaults to 0.1 s.
+    activation: nn.Module
+        Activation function. Defaults to ReLU activation.
+
+    References
+    ----------
+    .. [1] Zhang, C., Kim, Y. K., & Eskandarian, A. (2021).
+        EEG-inception: an accurate and robust end-to-end neural network
+        for EEG-based motor imagery classification.
+        Journal of Neural Engineering, 18(4), 046014.
+    """
+
+    def __init__(
+        self,
+        n_chans=None,
+        n_outputs=None,
+        input_window_seconds=None,
+        sfreq=250,
+        n_convs: int = 5,
+        n_filters: int = 48,
+        kernel_unit_s: float = 0.1,
+        activation: type[nn.Module] = nn.ReLU,
+        chs_info=None,
+        n_times=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
+
+        self.n_convs = n_convs
+        self.n_filters = n_filters
+        self.kernel_unit_s = kernel_unit_s
+        self.activation = activation
+
+        self.ensuredims = Ensure4d()
+        self.dimshuffle = Rearrange("batch C T 1 -> batch C 1 T")
+
+        self.mapping = {
+            "fc.weight": "final_layer.fc.weight",
+            "tc.bias": "final_layer.fc.bias",
+        }
+
+        # ======== Inception branches ========================
+
+        self.initial_inception_module = _InceptionModuleMI(
+            in_channels=self.n_chans,
+            n_filters=self.n_filters,
+            n_convs=self.n_convs,
+            kernel_unit_s=self.kernel_unit_s,
+            sfreq=self.sfreq,
+            activation=self.activation,
+        )
+
+        intermediate_in_channels = (self.n_convs + 1) * self.n_filters
+
+        self.intermediate_inception_modules_1 = nn.ModuleList(
+            [
+                _InceptionModuleMI(
+                    in_channels=intermediate_in_channels,
+                    n_filters=self.n_filters,
+                    n_convs=self.n_convs,
+                    kernel_unit_s=self.kernel_unit_s,
+                    sfreq=self.sfreq,
+                    activation=self.activation,
+                )
+                for _ in range(2)
+            ]
+        )
+
+        self.residual_block_1 = _ResidualModuleMI(
+            in_channels=self.n_chans,
+            n_filters=intermediate_in_channels,
+            activation=self.activation,
+        )
+
+        self.intermediate_inception_modules_2 = nn.ModuleList(
+            [
+                _InceptionModuleMI(
+                    in_channels=intermediate_in_channels,
+                    n_filters=self.n_filters,
+                    n_convs=self.n_convs,
+                    kernel_unit_s=self.kernel_unit_s,
+                    sfreq=self.sfreq,
+                    activation=self.activation,
+                )
+                for _ in range(3)
+            ]
+        )
+
+        self.residual_block_2 = _ResidualModuleMI(
+            in_channels=intermediate_in_channels,
+            n_filters=intermediate_in_channels,
+            activation=self.activation,
+        )
+
+        # XXX The paper mentions a final average pooling but does not indicate
+        # the kernel size... The only info available is figure1 showing a
+        # final AveragePooling layer and the table3 indicating the spatial and
+        # channel dimensions are unchanged by this layer... This could indicate
+        # a stride=1 as for MaxPooling layers. Howevere, when we look at the
+        # number of parameters of the linear layer following the average
+        # pooling, we see a small number of parameters, potentially indicating
+        # that the whole time dimension is averaged on this stage for each
+        # channel. We follow this last hypothesis here to comply with the
+        # number of parameters reported in the paper.
+        self.ave_pooling = nn.AvgPool2d(
+            kernel_size=(1, self.n_times),
+        )
+
+        self.flat = nn.Flatten()
+
+        module = nn.Sequential()
+        module.add_module(
+            "fc",
+            nn.Linear(
+                in_features=intermediate_in_channels,
+                out_features=self.n_outputs,
+                bias=True,
+            ),
+        )
+        module.add_module("out_fun", nn.Identity())
+        self.final_layer = module
+
+    def forward(
+        self,
+        X: torch.Tensor,
+    ) -> torch.Tensor:
+        X = self.ensuredims(X)
+        X = self.dimshuffle(X)
+
+        res1 = self.residual_block_1(X)
+
+        out = self.initial_inception_module(X)
+        for layer in self.intermediate_inception_modules_1:
+            out = layer(out)
+
+        out = out + res1
+
+        res2 = self.residual_block_2(out)
+
+        for layer in self.intermediate_inception_modules_2:
+            out = layer(out)
+
+        out = res2 + out
+
+        out = self.ave_pooling(out)
+        out = self.flat(out)
+        out = self.final_layer(out)
+        return out
+
+
+class _InceptionModuleMI(nn.Module):
+    r"""
+    Inception module.
+
+    This module implements a inception-like architecture that processes input
+    feature maps through multiple convolutional paths with different kernel
+    sizes, allowing the network to capture features at multiple scales.
+    It includes bottleneck layers, convolutional layers, and pooling layers,
+    followed by batch normalization and an activation function.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input channels in the input tensor.
+    n_filters : int
+        Number of filters (output channels) for each convolutional layer.
+    n_convs : int
+        Number of convolutional layers in the module (excluding bottleneck
+        and pooling paths).
+    kernel_unit_s : float, optional
+        Base size (in seconds) for the convolutional kernels. The actual kernel
+        size is computed as ``(2 * n_units + 1) * kernel_unit``, where ``n_units``
+        ranges from 0 to ``n_convs - 1``. Default is 0.1 seconds.
+    sfreq : float, optional
+        Sampling frequency of the input data, used to convert kernel sizes from
+        seconds to samples. Default is 250 Hz.
+    activation : nn.Module class, optional
+        Activation function class to apply after batch normalization. Should be
+        a PyTorch activation module class like ``nn.ReLU`` or ``nn.ELU``.
+        Default is ``nn.ReLU``.
+
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        n_filters,
+        n_convs,
+        kernel_unit_s=0.1,
+        sfreq=250,
+        activation: type[nn.Module] = nn.ReLU,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.n_filters = n_filters
+        self.n_convs = n_convs
+        self.kernel_unit_s = kernel_unit_s
+        self.sfreq = sfreq
+
+        self.bottleneck = nn.Conv2d(
+            in_channels=self.in_channels,
+            out_channels=self.n_filters,
+            kernel_size=1,
+            bias=True,
+        )
+
+        kernel_unit = int(self.kernel_unit_s * self.sfreq)
+
+        # XXX Maxpooling is usually used to reduce spatial resolution, with a
+        # stride equal to the kernel size... But it seems the authors use
+        # stride=1 in their paper according to the output shapes from Table3,
+        # although this is not clearly specified in the paper text.
+        self.pooling = nn.MaxPool2d(
+            kernel_size=(1, kernel_unit),
+            stride=1,
+            padding=(0, int(kernel_unit // 2)),
+        )
+
+        self.pooling_conv = nn.Conv2d(
+            in_channels=self.in_channels,
+            out_channels=self.n_filters,
+            kernel_size=1,
+            bias=True,
+        )
+
+        self.conv_list = nn.ModuleList(
+            [
+                nn.Conv2d(
+                    in_channels=self.n_filters,
+                    out_channels=self.n_filters,
+                    kernel_size=(1, (n_units * 2 + 1) * kernel_unit),
+                    padding="same",
+                    bias=True,
+                )
+                for n_units in range(self.n_convs)
+            ]
+        )
+
+        self.bn = nn.BatchNorm2d(self.n_filters * (self.n_convs + 1))
+
+        self.activation = activation()
+
+    def forward(
+        self,
+        X: torch.Tensor,
+    ) -> torch.Tensor:
+        X1 = self.bottleneck(X)
+
+        X1 = [conv(X1) for conv in self.conv_list]
+
+        X2 = self.pooling(X)
+        X2 = self.pooling_conv(X2)
+        # Get the target length from one of the conv branches
+        target_len = X1[0].shape[-1]
+
+        # Crop the pooling output if its length does not match
+        if X2.shape[-1] != target_len:
+            X2 = X2[..., :target_len]
+
+        out = torch.cat(X1 + [X2], 1)
+
+        out = self.bn(out)
+        return self.activation(out)
+
+
+class _ResidualModuleMI(nn.Module):
+    r"""
+    Residual module.
+
+    This module performs a 1x1 convolution followed by batch normalization and an activation function.
+    It is designed to process input feature maps and produce transformed output feature maps, often used
+    in residual connections within neural network architectures.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input channels in the input tensor.
+    n_filters : int
+        Number of filters (output channels) for the convolutional layer.
+    activation : nn.Module, optional
+        Activation function to apply after batch normalization. Should be an instance of a PyTorch
+        activation module (e.g., ``nn.ReLU()``, ``nn.ELU()``). Default is ``nn.ReLU()``.
+
+
+    """
+
+    def __init__(self, in_channels, n_filters, activation: type[nn.Module] = nn.ReLU):
+        super().__init__()
+        self.in_channels = in_channels
+        self.n_filters = n_filters
+        self.activation = activation()
+
+        self.bn = nn.BatchNorm2d(self.n_filters)
+        self.conv = nn.Conv2d(
+            in_channels=self.in_channels,
+            out_channels=self.n_filters,
+            kernel_size=1,
+            bias=True,
+        )
+
+    def forward(
+        self,
+        X: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Forward pass of the residual module.
+
+        Parameters
+        ----------
+        X : torch.Tensor
+            Input tensor of shape (batch_size, ch_names, n_times).
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor after convolution, batch normalization, and activation function.
+        """
+        out = self.conv(X)
+        out = self.bn(out)
+        return self.activation(out)

braindecode/models/eegitnet.py

@@ -0,0 +1,302 @@
+# Authors: Ghaith Bouallegue <ghaithbouallegue@gmail.com>
+#
+# License: BSD-3
+from einops.layers.torch import Rearrange
+from torch import nn
+
+from braindecode.models.base import EEGModuleMixin
+from braindecode.modules import DepthwiseConv2d, Ensure4d, InceptionBlock
+
+
+class EEGITNet(EEGModuleMixin, nn.Sequential):
+    r"""EEG-ITNet from Salami, et al (2022) [Salami2022]_
+
+    :bdg-success:`Convolution` :bdg-secondary:`Recurrent`
+
+    .. figure:: https://braindecode.org/dev/_static/model/eegitnet.jpg
+        :align: center
+        :alt: EEG-ITNet Architecture
+
+    EEG-ITNet: An Explainable Inception Temporal
+    Convolutional Network for motor imagery classification from
+    Salami et al. 2022.
+
+    See [Salami2022]_ for details.
+
+    Code adapted from https://github.com/abbassalami/eeg-itnet
+
+    Parameters
+    ----------
+    drop_prob: float
+        Dropout probability.
+    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``.
+    kernel_length : int, optional
+        Kernel length for inception branches. Determines the temporal receptive field.
+        Default is 16.
+    pool_kernel : int, optional
+        Pooling kernel size for the average pooling layer. Default is 4.
+    tcn_in_channel : int, optional
+        Number of input channels for Temporal Convolutional (TC) blocks. Default is 14.
+    tcn_kernel_size : int, optional
+        Kernel size for the TC blocks. Determines the temporal receptive field.
+        Default is 4.
+    tcn_padding : int, optional
+        Padding size for the TC blocks to maintain the input dimensions. Default is 3.
+    drop_prob : float, optional
+        Dropout probability applied after certain layers to prevent overfitting.
+        Default is 0.4.
+    tcn_dilatation : int, optional
+        Dilation rate for the first TC block. Subsequent blocks will have
+        dilation rates multiplied by powers of 2. Default is 1.
+
+    Notes
+    -----
+    This implementation is not guaranteed to be correct, has not been checked
+    by original authors, only reimplemented from the paper based on author implementation.
+
+
+    References
+    ----------
+    .. [Salami2022] A. Salami, J. Andreu-Perez and H. Gillmeister, "EEG-ITNet:
+        An Explainable Inception Temporal Convolutional Network for motor
+        imagery classification," in IEEE Access,
+        doi: 10.1109/ACCESS.2022.3161489.
+    """
+
+    def __init__(
+        self,
+        # Braindecode parameters
+        n_outputs=None,
+        n_chans=None,
+        n_times=None,
+        chs_info=None,
+        input_window_seconds=None,
+        sfreq=None,
+        # Model parameters
+        n_filters_time: int = 2,
+        kernel_length: int = 16,
+        pool_kernel: int = 4,
+        tcn_in_channel: int = 14,
+        tcn_kernel_size: int = 4,
+        tcn_padding: int = 3,
+        drop_prob: float = 0.4,
+        tcn_dilatation: int = 1,
+        activation: type[nn.Module] = nn.ELU,
+    ):
+        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,
+        )
+        self.mapping = {
+            "classification.1.weight": "final_layer.clf.weight",
+            "classification.1.bias": "final_layer.clf.weight",
+        }
+
+        del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
+
+        # ======== Handling EEG input ========================
+        self.add_module(
+            "input_preprocess",
+            nn.Sequential(Ensure4d(), Rearrange("ba ch t 1 -> ba 1 ch t")),
+        )
+        # ======== Inception branches ========================
+        block11 = self._get_inception_branch(
+            in_channels=self.n_chans,
+            out_channels=n_filters_time,
+            kernel_length=kernel_length,
+            activation=activation,
+        )
+        block12 = self._get_inception_branch(
+            in_channels=self.n_chans,
+            out_channels=n_filters_time * 2,
+            kernel_length=kernel_length * 2,
+            activation=activation,
+        )
+        block13 = self._get_inception_branch(
+            in_channels=self.n_chans,
+            out_channels=n_filters_time * 4,
+            kernel_length=n_filters_time * 4,
+            activation=activation,
+        )
+        self.add_module("inception_block", InceptionBlock((block11, block12, block13)))
+        self.pool1 = self.add_module(
+            "pooling",
+            nn.Sequential(
+                nn.AvgPool2d(kernel_size=(1, pool_kernel)), nn.Dropout(drop_prob)
+            ),
+        )
+        # =========== TC blocks =====================
+        self.add_module(
+            "TC_block1",
+            _TCBlock(
+                in_ch=tcn_in_channel,
+                kernel_length=tcn_kernel_size,
+                dilatation=tcn_dilatation,
+                padding=tcn_padding,
+                drop_prob=drop_prob,
+                activation=activation,
+            ),
+        )
+        # ================================
+        self.add_module(
+            "TC_block2",
+            _TCBlock(
+                in_ch=tcn_in_channel,
+                kernel_length=tcn_kernel_size,
+                dilatation=tcn_dilatation * 2,
+                padding=tcn_padding * 2,
+                drop_prob=drop_prob,
+                activation=activation,
+            ),
+        )
+        # ================================
+        self.add_module(
+            "TC_block3",
+            _TCBlock(
+                in_ch=tcn_in_channel,
+                kernel_length=tcn_kernel_size,
+                dilatation=tcn_dilatation * 4,
+                padding=tcn_padding * 4,
+                drop_prob=drop_prob,
+                activation=activation,
+            ),
+        )
+        # ================================
+        self.add_module(
+            "TC_block4",
+            _TCBlock(
+                in_ch=tcn_in_channel,
+                kernel_length=tcn_kernel_size,
+                dilatation=tcn_dilatation * 8,
+                padding=tcn_padding * 8,
+                drop_prob=drop_prob,
+                activation=activation,
+            ),
+        )
+
+        # ============= Dimensionality reduction ===================
+        self.add_module(
+            "dim_reduction",
+            nn.Sequential(
+                nn.Conv2d(tcn_in_channel, tcn_in_channel * 2, kernel_size=(1, 1)),
+                nn.BatchNorm2d(tcn_in_channel * 2),
+                activation(),
+                nn.AvgPool2d((1, tcn_kernel_size)),
+                nn.Dropout(drop_prob),
+            ),
+        )
+        # ============== Classifier ==================
+        # Moved flatten to another layer
+        self.add_module("flatten", nn.Flatten())
+
+        num_features = self.get_output_shape()[-1]
+
+        self.add_module("final_layer", nn.Linear(num_features, self.n_outputs))
+
+    @staticmethod
+    def _get_inception_branch(
+        in_channels,
+        out_channels,
+        kernel_length,
+        depth_multiplier=1,
+        activation: type[nn.Module] = nn.ELU,
+    ):
+        return nn.Sequential(
+            nn.Conv2d(
+                1,
+                out_channels,
+                kernel_size=(1, kernel_length),
+                padding="same",
+                bias=False,
+            ),
+            nn.BatchNorm2d(out_channels),
+            DepthwiseConv2d(
+                out_channels,
+                kernel_size=(in_channels, 1),
+                depth_multiplier=depth_multiplier,
+                bias=False,
+                padding="valid",
+            ),
+            nn.BatchNorm2d(out_channels),
+            activation(),
+        )
+
+
+class _TCBlock(nn.Module):
+    r"""
+    Temporal Convolutional (TC) block.
+
+    This module applies two depthwise separable convolutions with dilation and residual
+    connections, commonly used in temporal convolutional networks to capture long-range
+    dependencies in time-series data.
+
+    Parameters
+    ----------
+    in_ch : int
+        Number of input channels.
+    kernel_length : int
+        Length of the convolutional kernels.
+    dilatation : int
+        Dilatation rate for the convolutions.
+    padding : int
+        Amount of padding to add to the input.
+    drop_prob : float, optional
+        Dropout probability. Default is 0.4.
+    activation : nn.Module class, optional
+        Activation function class to use. Should be a PyTorch activation module class
+        like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ELU``.
+    """
+
+    def __init__(
+        self,
+        in_ch,
+        kernel_length,
+        dilatation,
+        padding,
+        drop_prob=0.4,
+        activation: type[nn.Module] = nn.ELU,
+    ):
+        super().__init__()
+        self.pad = padding
+        self.tc1 = nn.Sequential(
+            DepthwiseConv2d(
+                in_ch,
+                kernel_size=(1, kernel_length),
+                depth_multiplier=1,
+                dilation=(1, dilatation),
+                bias=False,
+                padding="valid",
+            ),
+            nn.BatchNorm2d(in_ch),
+            activation(),
+            nn.Dropout(drop_prob),
+        )
+
+        self.tc2 = nn.Sequential(
+            DepthwiseConv2d(
+                in_ch,
+                kernel_size=(1, kernel_length),
+                depth_multiplier=1,
+                dilation=(1, dilatation),
+                bias=False,
+                padding="valid",
+            ),
+            nn.BatchNorm2d(in_ch),
+            activation(),
+            nn.Dropout(drop_prob),
+        )
+
+    def forward(self, x):
+        residual = x
+        paddings = (self.pad, 0, 0, 0, 0, 0, 0, 0)
+        x = nn.functional.pad(x, paddings)
+        x = self.tc1(x)
+        x = nn.functional.pad(x, paddings)
+        x = self.tc2(x) + residual
+        return x