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

braindecode/models/eegconformer.py

@@ -0,0 +1,372 @@
+# Authors: Yonghao Song <eeyhsong@gmail.com>
+#
+# License: BSD (3-clause)
+import warnings
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from einops.layers.torch import Rearrange
+from torch import Tensor, nn
+
+from braindecode.models.base import EEGModuleMixin
+from braindecode.modules import FeedForwardBlock, MultiHeadAttention
+
+
+class EEGConformer(EEGModuleMixin, nn.Module):
+    """EEG Conformer from Song et al. (2022) from [song2022]_.
+
+     .. figure:: https://raw.githubusercontent.com/eeyhsong/EEG-Conformer/refs/heads/main/visualization/Fig1.png
+        :align: center
+        :alt: EEGConformer Architecture
+
+    Convolutional Transformer for EEG decoding.
+
+    The paper and original code with more details about the methodological
+    choices are available at the [song2022]_ and [ConformerCode]_.
+
+    This neural network architecture receives a traditional braindecode input.
+    The input shape should be three-dimensional matrix representing the EEG
+    signals.
+
+         `(batch_size, n_channels, n_timesteps)`.
+
+    The EEG Conformer architecture is composed of three modules:
+        - PatchEmbedding
+        - TransformerEncoder
+        - ClassificationHead
+
+    Notes
+    -----
+    The authors recommend using data augmentation before using Conformer,
+    e.g. segmentation and recombination,
+    Please refer to the original paper and code for more details.
+
+    The model was initially tuned on 4 seconds of 250 Hz data.
+    Please adjust the scale of the temporal convolutional layer,
+    and the pooling layer for better performance.
+
+    .. versionadded:: 0.8
+
+    We aggregate the parameters based on the parts of the models, or
+    when the parameters were used first, e.g. n_filters_time.
+
+    Parameters
+    ----------
+    n_filters_time: int
+        Number of temporal filters, defines also embedding size.
+    filter_time_length: int
+        Length of the temporal filter.
+    pool_time_length: int
+        Length of temporal pooling filter.
+    pool_time_stride: int
+        Length of stride between temporal pooling filters.
+    drop_prob: float
+        Dropout rate of the convolutional layer.
+    att_depth: int
+        Number of self-attention layers.
+    att_heads: int
+        Number of attention heads.
+    att_drop_prob: float
+        Dropout rate of the self-attention layer.
+    final_fc_length: int | str
+        The dimension of the fully connected layer.
+    return_features: bool
+        If True, the forward method returns the features before the
+        last classification layer. Defaults to False.
+    activation: nn.Module
+        Activation function as parameter. Default is nn.ELU
+    activation_transfor: nn.Module
+        Activation function as parameter, applied at the FeedForwardBlock module
+        inside the transformer. Default is nn.GeLU
+
+    References
+    ----------
+    .. [song2022] Song, Y., Zheng, Q., Liu, B. and Gao, X., 2022. EEG
+       conformer: Convolutional transformer for EEG decoding and visualization.
+       IEEE Transactions on Neural Systems and Rehabilitation Engineering,
+       31, pp.710-719. https://ieeexplore.ieee.org/document/9991178
+    .. [ConformerCode] Song, Y., Zheng, Q., Liu, B. and Gao, X., 2022. EEG
+       conformer: Convolutional transformer for EEG decoding and visualization.
+       https://github.com/eeyhsong/EEG-Conformer.
+    """
+
+    def __init__(
+        self,
+        n_outputs=None,
+        n_chans=None,
+        n_filters_time=40,
+        filter_time_length=25,
+        pool_time_length=75,
+        pool_time_stride=15,
+        drop_prob=0.5,
+        att_depth=6,
+        att_heads=10,
+        att_drop_prob=0.5,
+        final_fc_length="auto",
+        return_features=False,
+        activation: nn.Module = nn.ELU,
+        activation_transfor: nn.Module = nn.GELU,
+        n_times=None,
+        chs_info=None,
+        input_window_seconds=None,
+        sfreq=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,
+        )
+        self.mapping = {
+            "classification_head.fc.6.weight": "final_layer.final_layer.0.weight",
+            "classification_head.fc.6.bias": "final_layer.final_layer.0.bias",
+        }
+
+        del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
+        if not (self.n_chans <= 64):
+            warnings.warn(
+                "This model has only been tested on no more "
+                + "than 64 channels. no guarantee to work with "
+                + "more channels.",
+                UserWarning,
+            )
+
+        self.return_features = return_features
+
+        self.patch_embedding = _PatchEmbedding(
+            n_filters_time=n_filters_time,
+            filter_time_length=filter_time_length,
+            n_channels=self.n_chans,
+            pool_time_length=pool_time_length,
+            stride_avg_pool=pool_time_stride,
+            drop_prob=drop_prob,
+            activation=activation,
+        )
+
+        if final_fc_length == "auto":
+            assert self.n_times is not None
+            self.final_fc_length = self.get_fc_size()
+
+        self.transformer = _TransformerEncoder(
+            att_depth=att_depth,
+            emb_size=n_filters_time,
+            att_heads=att_heads,
+            att_drop=att_drop_prob,
+            activation=activation_transfor,
+        )
+
+        self.fc = _FullyConnected(
+            final_fc_length=self.final_fc_length, activation=activation
+        )
+
+        self.final_layer = nn.Linear(self.fc.hidden_channels, self.n_outputs)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = torch.unsqueeze(x, dim=1)  # add one extra dimension
+        x = self.patch_embedding(x)
+        feature = self.transformer(x)
+
+        if self.return_features:
+            return feature
+
+        x = self.fc(feature)
+        x = self.final_layer(x)
+        return x
+
+    def get_fc_size(self):
+        out = self.patch_embedding(torch.ones((1, 1, self.n_chans, self.n_times)))
+        size_embedding_1 = out.cpu().data.numpy().shape[1]
+        size_embedding_2 = out.cpu().data.numpy().shape[2]
+
+        return size_embedding_1 * size_embedding_2
+
+
+class _PatchEmbedding(nn.Module):
+    """Patch Embedding.
+
+    The authors used a convolution module to capture local features,
+    instead of position embedding.
+
+    Parameters
+    ----------
+    n_filters_time: int
+        Number of temporal filters, defines also embedding size.
+    filter_time_length: int
+        Length of the temporal filter.
+    n_channels: int
+        Number of channels to be used as number of spatial filters.
+    pool_time_length: int
+        Length of temporal poling filter.
+    stride_avg_pool: int
+        Length of stride between temporal pooling filters.
+    drop_prob: float
+        Dropout rate of the convolutional layer.
+
+    Returns
+    -------
+    x: torch.Tensor
+        The output tensor of the patch embedding layer.
+    """
+
+    def __init__(
+        self,
+        n_filters_time,
+        filter_time_length,
+        n_channels,
+        pool_time_length,
+        stride_avg_pool,
+        drop_prob,
+        activation: nn.Module = nn.ELU,
+    ):
+        super().__init__()
+
+        self.shallownet = nn.Sequential(
+            nn.Conv2d(1, n_filters_time, (1, filter_time_length), (1, 1)),
+            nn.Conv2d(n_filters_time, n_filters_time, (n_channels, 1), (1, 1)),
+            nn.BatchNorm2d(num_features=n_filters_time),
+            activation(),
+            nn.AvgPool2d(
+                kernel_size=(1, pool_time_length), stride=(1, stride_avg_pool)
+            ),
+            # pooling acts as slicing to obtain 'patch' along the
+            # time dimension as in ViT
+            nn.Dropout(p=drop_prob),
+        )
+
+        self.projection = nn.Sequential(
+            nn.Conv2d(
+                n_filters_time, n_filters_time, (1, 1), stride=(1, 1)
+            ),  # transpose, conv could enhance fiting ability slightly
+            Rearrange("b d_model 1 seq -> b seq d_model"),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.shallownet(x)
+        x = self.projection(x)
+        return x
+
+
+class _ResidualAdd(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x):
+        res = x
+        x = self.fn(x)
+        x += res
+        return x
+
+
+class _TransformerEncoderBlock(nn.Sequential):
+    def __init__(
+        self,
+        emb_size,
+        att_heads,
+        att_drop,
+        forward_expansion=4,
+        activation: nn.Module = nn.GELU,
+    ):
+        super().__init__(
+            _ResidualAdd(
+                nn.Sequential(
+                    nn.LayerNorm(emb_size),
+                    MultiHeadAttention(emb_size, att_heads, att_drop),
+                    nn.Dropout(att_drop),
+                )
+            ),
+            _ResidualAdd(
+                nn.Sequential(
+                    nn.LayerNorm(emb_size),
+                    FeedForwardBlock(
+                        emb_size,
+                        expansion=forward_expansion,
+                        drop_p=att_drop,
+                        activation=activation,
+                    ),
+                    nn.Dropout(att_drop),
+                )
+            ),
+        )
+
+
+class _TransformerEncoder(nn.Sequential):
+    """Transformer encoder module for the transformer encoder.
+
+    Similar to the layers used in ViT.
+
+    Parameters
+    ----------
+    att_depth : int
+        Number of transformer encoder blocks.
+    emb_size : int
+        Embedding size of the transformer encoder.
+    att_heads : int
+        Number of attention heads.
+    att_drop : float
+        Dropout probability for the attention layers.
+
+    """
+
+    def __init__(
+        self, att_depth, emb_size, att_heads, att_drop, activation: nn.Module = nn.GELU
+    ):
+        super().__init__(
+            *[
+                _TransformerEncoderBlock(
+                    emb_size, att_heads, att_drop, activation=activation
+                )
+                for _ in range(att_depth)
+            ]
+        )
+
+
+class _FullyConnected(nn.Module):
+    def __init__(
+        self,
+        final_fc_length,
+        drop_prob_1=0.5,
+        drop_prob_2=0.3,
+        out_channels=256,
+        hidden_channels=32,
+        activation: nn.Module = nn.ELU,
+    ):
+        """Fully-connected layer for the transformer encoder.
+
+        Parameters
+        ----------
+        final_fc_length : int
+            Length of the final fully connected layer.
+        n_classes : int
+            Number of classes for classification.
+        drop_prob_1 : float
+            Dropout probability for the first dropout layer.
+        drop_prob_2 : float
+            Dropout probability for the second dropout layer.
+        out_channels : int
+            Number of output channels for the first linear layer.
+        hidden_channels : int
+            Number of output channels for the second linear layer.
+        return_features : bool
+            Whether to return input features.
+        """
+
+        super().__init__()
+        self.hidden_channels = hidden_channels
+        self.fc = nn.Sequential(
+            nn.Linear(final_fc_length, out_channels),
+            activation(),
+            nn.Dropout(drop_prob_1),
+            nn.Linear(out_channels, hidden_channels),
+            activation(),
+            nn.Dropout(drop_prob_2),
+        )
+
+    def forward(self, x):
+        x = x.contiguous().view(x.size(0), -1)
+        out = self.fc(x)
+        return out

braindecode/models/eeginception_erp.py

@@ -0,0 +1,304 @@
+# Authors: Bruno Aristimunha <b.aristimunha@gmail.com>
+#          Cedric Rommel <cedric.rommel@inria.fr>
+#
+# License: BSD (3-clause)
+import math
+
+from einops.layers.torch import Rearrange
+from torch import nn
+
+from braindecode.functional import glorot_weight_zero_bias
+from braindecode.models.base import EEGModuleMixin
+from braindecode.modules import DepthwiseConv2d, Ensure4d, InceptionBlock
+
+
+class EEGInceptionERP(EEGModuleMixin, nn.Sequential):
+    """EEG Inception for ERP-based from Santamaria-Vazquez et al (2020) [santamaria2020]_.
+
+    .. figure:: https://braindecode.org/dev/_static/model/eeginceptionerp.jpg
+        :align: center
+        :alt: EEGInceptionERP Architecture
+
+    The code for the paper and this model is also available at [santamaria2020]_
+    and an adaptation for PyTorch [2]_.
+
+    The model is strongly based on the original InceptionNet for an image. The main goal is
+    to extract features in parallel with different scales. The authors extracted three scales
+    proportional to the window sample size. The network had three parts:
+    1-larger inception block largest, 2-smaller inception block followed by 3-bottleneck
+    for classification.
+
+    One advantage of the EEG-Inception block is that it allows a network
+    to learn simultaneous components of low and high frequency associated with the signal.
+    The winners of BEETL Competition/NeurIps 2021 used parts of the
+    model [beetl]_.
+
+    The model is fully described in [santamaria2020]_.
+
+    Notes
+    -----
+    This implementation is not guaranteed to be correct, has not been checked
+    by original authors, only reimplemented from the paper based on [2]_.
+
+    Parameters
+    ----------
+    n_times : int, optional
+        Size of the input, in number of samples. Set to 128 (1s) as in
+        [santamaria2020]_.
+    sfreq : float, optional
+        EEG sampling frequency. Defaults to 128 as in [santamaria2020]_.
+    drop_prob : float, optional
+        Dropout rate inside all the network. Defaults to 0.5 as in
+        [santamaria2020]_.
+    scales_samples_s: list(float), optional
+        Windows for inception block. Temporal scale (s) of the convolutions on
+        each Inception module. This parameter determines the kernel sizes of
+        the filters. Defaults to 0.5, 0.25, 0.125 seconds, as in
+        [santamaria2020]_.
+    n_filters : int, optional
+        Initial number of convolutional filters. Defaults to 8 as in
+        [santamaria2020]_.
+    activation: nn.Module, optional
+        Activation function. Defaults to ELU activation as in
+        [santamaria2020]_.
+    batch_norm_alpha: float, optional
+        Momentum for BatchNorm2d. Defaults to 0.01.
+    depth_multiplier: int, optional
+        Depth multiplier for the depthwise convolution. Defaults to 2 as in
+        [santamaria2020]_.
+    pooling_sizes: list(int), optional
+        Pooling sizes for the inception blocks. Defaults to 4, 2, 2 and 2, as
+        in [santamaria2020]_.
+
+
+    References
+    ----------
+    .. [santamaria2020] Santamaria-Vazquez, E., Martinez-Cagigal, V.,
+       Vaquerizo-Villar, F., & Hornero, R. (2020).
+       EEG-inception: A novel deep convolutional neural network for assistive
+       ERP-based brain-computer interfaces.
+       IEEE Transactions on Neural Systems and Rehabilitation Engineering , v. 28.
+       Online: http://dx.doi.org/10.1109/TNSRE.2020.3048106
+    .. [2] Grifcc. Implementation of the EEGInception in torch (2022).
+       Online: https://github.com/Grifcc/EEG/
+    .. [beetl] Wei, X., Faisal, A.A., Grosse-Wentrup, M., Gramfort, A., Chevallier, S.,
+       Jayaram, V., Jeunet, C., Bakas, S., Ludwig, S., Barmpas, K., Bahri, M., Panagakis,
+       Y., Laskaris, N., Adamos, D.A., Zafeiriou, S., Duong, W.C., Gordon, S.M.,
+       Lawhern, V.J., Śliwowski, M., Rouanne, V. &amp; Tempczyk, P. (2022).
+       2021 BEETL Competition: Advancing Transfer Learning for Subject Independence &amp;
+       Heterogeneous EEG Data Sets. Proceedings of the NeurIPS 2021 Competitions and
+       Demonstrations Track, in Proceedings of Machine Learning Research
+       176:205-219 Available from https://proceedings.mlr.press/v176/wei22a.html.
+
+    """
+
+    def __init__(
+        self,
+        n_chans=None,
+        n_outputs=None,
+        n_times=1000,
+        sfreq=128,
+        drop_prob=0.5,
+        scales_samples_s=(0.5, 0.25, 0.125),
+        n_filters=8,
+        activation: nn.Module = nn.ELU,
+        batch_norm_alpha=0.01,
+        depth_multiplier=2,
+        pooling_sizes=(4, 2, 2, 2),
+        chs_info=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, input_window_seconds, sfreq
+        self.drop_prob = drop_prob
+        self.n_filters = n_filters
+        self.scales_samples_s = scales_samples_s
+        self.scales_samples = tuple(
+            int(size_s * self.sfreq) for size_s in self.scales_samples_s
+        )
+        self.activation = activation
+        self.alpha_momentum = batch_norm_alpha
+        self.depth_multiplier = depth_multiplier
+        self.pooling_sizes = pooling_sizes
+
+        self.mapping = {
+            "classification.1.weight": "final_layer.fc.weight",
+            "classification.1.bias": "final_layer.fc.bias",
+        }
+
+        self.add_module("ensuredims", Ensure4d())
+
+        self.add_module("dimshuffle", Rearrange("batch C T 1 -> batch 1 C T"))
+
+        # ======== Inception branches ========================
+        block11 = self._get_inception_branch_1(
+            in_channels=self.n_chans,
+            out_channels=self.n_filters,
+            kernel_length=self.scales_samples[0],
+            alpha_momentum=self.alpha_momentum,
+            activation=self.activation,
+            drop_prob=self.drop_prob,
+            depth_multiplier=self.depth_multiplier,
+        )
+        block12 = self._get_inception_branch_1(
+            in_channels=self.n_chans,
+            out_channels=self.n_filters,
+            kernel_length=self.scales_samples[1],
+            alpha_momentum=self.alpha_momentum,
+            activation=self.activation,
+            drop_prob=self.drop_prob,
+            depth_multiplier=self.depth_multiplier,
+        )
+        block13 = self._get_inception_branch_1(
+            in_channels=self.n_chans,
+            out_channels=self.n_filters,
+            kernel_length=self.scales_samples[2],
+            alpha_momentum=self.alpha_momentum,
+            activation=self.activation,
+            drop_prob=self.drop_prob,
+            depth_multiplier=self.depth_multiplier,
+        )
+
+        self.add_module(
+            "inception_block_1", InceptionBlock((block11, block12, block13))
+        )
+
+        self.add_module("avg_pool_1", nn.AvgPool2d((1, self.pooling_sizes[0])))
+
+        # ======== Inception branches ========================
+        n_concat_filters = len(self.scales_samples) * self.n_filters
+        n_concat_dw_filters = n_concat_filters * self.depth_multiplier
+        block21 = self._get_inception_branch_2(
+            in_channels=n_concat_dw_filters,
+            out_channels=self.n_filters,
+            kernel_length=self.scales_samples[0] // 4,
+            alpha_momentum=self.alpha_momentum,
+            activation=self.activation,
+            drop_prob=self.drop_prob,
+        )
+        block22 = self._get_inception_branch_2(
+            in_channels=n_concat_dw_filters,
+            out_channels=self.n_filters,
+            kernel_length=self.scales_samples[1] // 4,
+            alpha_momentum=self.alpha_momentum,
+            activation=self.activation,
+            drop_prob=self.drop_prob,
+        )
+        block23 = self._get_inception_branch_2(
+            in_channels=n_concat_dw_filters,
+            out_channels=self.n_filters,
+            kernel_length=self.scales_samples[2] // 4,
+            alpha_momentum=self.alpha_momentum,
+            activation=self.activation,
+            drop_prob=self.drop_prob,
+        )
+
+        self.add_module(
+            "inception_block_2", InceptionBlock((block21, block22, block23))
+        )
+
+        self.add_module("avg_pool_2", nn.AvgPool2d((1, self.pooling_sizes[1])))
+
+        self.add_module(
+            "final_block",
+            nn.Sequential(
+                nn.Conv2d(
+                    n_concat_filters,
+                    n_concat_filters // 2,
+                    (1, 8),
+                    padding="same",
+                    bias=False,
+                ),
+                nn.BatchNorm2d(n_concat_filters // 2, momentum=self.alpha_momentum),
+                activation(),
+                nn.Dropout(self.drop_prob),
+                nn.AvgPool2d((1, self.pooling_sizes[2])),
+                nn.Conv2d(
+                    n_concat_filters // 2,
+                    n_concat_filters // 4,
+                    (1, 4),
+                    padding="same",
+                    bias=False,
+                ),
+                nn.BatchNorm2d(n_concat_filters // 4, momentum=self.alpha_momentum),
+                activation(),
+                nn.Dropout(self.drop_prob),
+                nn.AvgPool2d((1, self.pooling_sizes[3])),
+            ),
+        )
+
+        spatial_dim_last_layer = self.n_times // math.prod(self.pooling_sizes)
+        n_channels_last_layer = self.n_filters * len(self.scales_samples) // 4
+
+        self.add_module("flat", nn.Flatten())
+
+        # Incorporating classification module and subsequent ones in one final layer
+        module = nn.Sequential()
+
+        module.add_module(
+            "fc",
+            nn.Linear(spatial_dim_last_layer * n_channels_last_layer, self.n_outputs),
+        )
+
+        module.add_module("identity", nn.Identity())
+
+        self.add_module("final_layer", module)
+
+        glorot_weight_zero_bias(self)
+
+    @staticmethod
+    def _get_inception_branch_1(
+        in_channels,
+        out_channels,
+        kernel_length,
+        alpha_momentum,
+        drop_prob,
+        activation,
+        depth_multiplier,
+    ):
+        return nn.Sequential(
+            nn.Conv2d(
+                1,
+                out_channels,
+                kernel_size=(1, kernel_length),
+                padding="same",
+                bias=True,
+            ),
+            nn.BatchNorm2d(out_channels, momentum=alpha_momentum),
+            activation(),
+            nn.Dropout(drop_prob),
+            DepthwiseConv2d(
+                out_channels,
+                kernel_size=(in_channels, 1),
+                depth_multiplier=depth_multiplier,
+                bias=False,
+                padding="valid",
+            ),
+            nn.BatchNorm2d(depth_multiplier * out_channels, momentum=alpha_momentum),
+            activation(),
+            nn.Dropout(drop_prob),
+        )
+
+    @staticmethod
+    def _get_inception_branch_2(
+        in_channels, out_channels, kernel_length, alpha_momentum, drop_prob, activation
+    ):
+        return nn.Sequential(
+            nn.Conv2d(
+                in_channels,
+                out_channels,
+                kernel_size=(1, kernel_length),
+                padding="same",
+                bias=False,
+            ),
+            nn.BatchNorm2d(out_channels, momentum=alpha_momentum),
+            activation(),
+            nn.Dropout(drop_prob),
+        )