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

braindecode/models/msvtnet.py

@@ -0,0 +1,375 @@
+# Authors: Tao Yang <sheeptao@outlook.com>
+#          Bruno Aristimunha <b.aristimunha@gmail.com> (braindecode adaptation)
+#
+from typing import Type, Union
+
+import torch
+import torch.nn as nn
+from einops.layers.torch import Rearrange
+
+from braindecode.models.base import EEGModuleMixin
+
+
+class MSVTNet(EEGModuleMixin, nn.Module):
+    """MSVTNet model from Liu K et al (2024) from [msvt2024]_.
+
+    This model implements a multi-scale convolutional transformer network
+    for EEG signal classification, as described in [msvt2024]_.
+
+    .. figure:: https://raw.githubusercontent.com/SheepTAO/MSVTNet/refs/heads/main/MSVTNet_Arch.png
+       :align: center
+       :alt: MSVTNet Architecture
+
+    Parameters
+    ----------
+    n_filters_list : list[int], optional
+        List of filter numbers for each TSConv block, by default (9, 9, 9, 9).
+    conv1_kernels_size : list[int], optional
+        List of kernel sizes for the first convolution in each TSConv block,
+        by default (15, 31, 63, 125).
+    conv2_kernel_size : int, optional
+        Kernel size for the second convolution in TSConv blocks, by default 15.
+    depth_multiplier : int, optional
+        Depth multiplier for depthwise convolution, by default 2.
+    pool1_size : int, optional
+        Pooling size for the first pooling layer in TSConv blocks, by default 8.
+    pool2_size : int, optional
+        Pooling size for the second pooling layer in TSConv blocks, by default 7.
+    drop_prob : float, optional
+        Dropout probability for convolutional layers, by default 0.3.
+    num_heads : int, optional
+        Number of attention heads in the transformer encoder, by default 8.
+    feedforward_ratio : float, optional
+        Ratio to compute feedforward dimension in the transformer, by default 1.
+    drop_prob_trans : float, optional
+        Dropout probability for the transformer, by default 0.5.
+    num_layers : int, optional
+        Number of transformer encoder layers, by default 2.
+    activation : Type[nn.Module], optional
+        Activation function class to use, by default nn.ELU.
+    return_features : bool, optional
+        Whether to return predictions from branch classifiers, by default False.
+
+    Notes
+    -----
+    This implementation is not guaranteed to be correct, has not been checked
+    by original authors, only reimplemented based on the original code [msvt2024code]_.
+
+    References
+    ----------
+    .. [msvt2024] Liu, K., et al. (2024). MSVTNet: Multi-Scale Vision
+       Transformer Neural Network for EEG-Based Motor Imagery Decoding.
+       IEEE Journal of Biomedical an Health Informatics.
+    .. [msvt2024code] Liu, K., et al. (2024). MSVTNet: Multi-Scale Vision
+       Transformer Neural Network for EEG-Based Motor Imagery Decoding.
+       Source Code: https://github.com/SheepTAO/MSVTNet
+    """
+
+    def __init__(
+        self,
+        # braindecode parameters
+        n_chans=None,
+        n_outputs=None,
+        n_times=None,
+        input_window_seconds=None,
+        sfreq=None,
+        chs_info=None,
+        # Model's parameters
+        n_filters_list: tuple[int, ...] = (9, 9, 9, 9),
+        conv1_kernels_size: tuple[int, ...] = (15, 31, 63, 125),
+        conv2_kernel_size: int = 15,
+        depth_multiplier: int = 2,
+        pool1_size: int = 8,
+        pool2_size: int = 7,
+        drop_prob: float = 0.3,
+        num_heads: int = 8,
+        feedforward_ratio: float = 1,
+        drop_prob_trans: float = 0.5,
+        num_layers: int = 2,
+        activation: Type[nn.Module] = nn.ELU,
+        return_features: bool = False,
+    ):
+        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.return_features = return_features
+        assert len(n_filters_list) == len(conv1_kernels_size), (
+            "The length of n_filters_list and conv1_kernel_sizes should be equal."
+        )
+
+        self.ensure_dim = Rearrange("batch chans time -> batch 1 chans time")
+        self.mstsconv = nn.ModuleList(
+            [
+                nn.Sequential(
+                    _TSConv(
+                        self.n_chans,
+                        n_filters_list[b],
+                        conv1_kernels_size[b],
+                        conv2_kernel_size,
+                        depth_multiplier,
+                        pool1_size,
+                        pool2_size,
+                        drop_prob,
+                        activation,
+                    ),
+                    Rearrange("batch channels 1 time -> batch time channels"),
+                )
+                for b in range(len(n_filters_list))
+            ]
+        )
+        branch_linear_in = self._forward_flatten(cat=False)
+        self.branch_head = nn.ModuleList(
+            [
+                _DenseLayers(branch_linear_in[b].shape[1], self.n_outputs)
+                for b in range(len(n_filters_list))
+            ]
+        )
+
+        seq_len, d_model = self._forward_mstsconv().shape[1:3]  # type: ignore
+        self.transformer = _Transformer(
+            seq_len,
+            d_model,
+            num_heads,
+            feedforward_ratio,
+            drop_prob_trans,
+            num_layers,
+        )
+
+        linear_in = self._forward_flatten().shape[1]  # type: ignore
+        self.flatten_layer = nn.Flatten()
+        self.final_layer = nn.Linear(linear_in, self.n_outputs)
+
+    def _forward_mstsconv(
+        self, cat: bool = True
+    ) -> Union[torch.Tensor, list[torch.Tensor]]:
+        x = torch.randn(1, 1, self.n_chans, self.n_times)
+        x = [tsconv(x) for tsconv in self.mstsconv]
+        if cat:
+            x = torch.cat(x, dim=2)
+        return x
+
+    def _forward_flatten(
+        self, cat: bool = True
+    ) -> Union[torch.Tensor, list[torch.Tensor]]:
+        x = self._forward_mstsconv(cat)
+        if cat:
+            x = self.transformer(x)
+            x = x.flatten(start_dim=1, end_dim=-1)
+        else:
+            x = [_.flatten(start_dim=1, end_dim=-1) for _ in x]
+        return x
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x with shape: (batch, n_chans, n_times)
+        x = self.ensure_dim(x)
+        # x with shape: (batch, 1, n_chans, n_times)
+        x_list = [tsconv(x) for tsconv in self.mstsconv]
+        # x_list contains 4 tensors, each of shape: [batch_size, seq_len, embed_dim]
+        branch_preds = [
+            branch(x_list[idx]) for idx, branch in enumerate(self.branch_head)
+        ]
+        # branch_preds contains 4 tensors, each of shape: [batch_size, num_classes]
+        x = torch.stack(x_list, dim=2)
+        x = x.view(x.size(0), x.size(1), -1)
+        # x shape after concatenation: [batch_size, seq_len, total_embed_dim]
+        x = self.transformer(x)
+        # x shape after transformer: [batch_size, embed_dim]
+
+        x = self.final_layer(x)
+        if self.return_features:
+            # x shape after final layer: [batch_size, num_classes]
+            # branch_preds shape: [batch_size, num_classes]
+            return torch.stack(branch_preds)
+        return x
+
+
+class _TSConv(nn.Sequential):
+    """
+    Time-Distributed Separable Convolution block.
+
+    The architecture consists of:
+    - **Temporal Convolution**
+    - **Batch Normalization**
+    - **Depthwise Spatial Convolution**
+    - **Batch Normalization**
+    - **Activation Function**
+    - **First Pooling Layer**
+    - **Dropout**
+    - **Depthwise Temporal Convolution**
+    - **Batch Normalization**
+    - **Activation Function**
+    - **Second Pooling Layer**
+    - **Dropout**
+
+    Parameters
+    ----------
+    n_channels : int
+        Number of input channels (EEG channels).
+    n_filters : int
+        Number of filters for the convolution layers.
+    conv1_kernel_size : int
+        Kernel size for the first convolution layer.
+    conv2_kernel_size : int
+        Kernel size for the second convolution layer.
+    depth_multiplier : int
+        Depth multiplier for depthwise convolution.
+    pool1_size : int
+        Kernel size for the first pooling layer.
+    pool2_size : int
+        Kernel size for the second pooling layer.
+    drop_prob : float
+        Dropout probability.
+    activation : Type[nn.Module], optional
+        Activation function class to use, by default nn.ELU.
+    """
+
+    def __init__(
+        self,
+        n_channels: int,
+        n_filters: int,
+        conv1_kernel_size: int,
+        conv2_kernel_size: int,
+        depth_multiplier: int,
+        pool1_size: int,
+        pool2_size: int,
+        drop_prob: float,
+        activation: Type[nn.Module] = nn.ELU,
+    ):
+        super().__init__(
+            nn.Conv2d(
+                in_channels=1,
+                out_channels=n_filters,
+                kernel_size=(1, conv1_kernel_size),
+                padding="same",
+                bias=False,
+            ),
+            nn.BatchNorm2d(n_filters),
+            nn.Conv2d(
+                in_channels=n_filters,
+                out_channels=n_filters * depth_multiplier,
+                kernel_size=(n_channels, 1),
+                groups=n_filters,
+                bias=False,
+            ),
+            nn.BatchNorm2d(n_filters * depth_multiplier),
+            activation(),
+            nn.AvgPool2d(kernel_size=(1, pool1_size)),
+            nn.Dropout(drop_prob),
+            nn.Conv2d(
+                in_channels=n_filters * depth_multiplier,
+                out_channels=n_filters * depth_multiplier,
+                kernel_size=(1, conv2_kernel_size),
+                padding="same",
+                groups=n_filters * depth_multiplier,
+                bias=False,
+            ),
+            nn.BatchNorm2d(n_filters * depth_multiplier),
+            activation(),
+            nn.AvgPool2d(kernel_size=(1, pool2_size)),
+            nn.Dropout(drop_prob),
+        )
+
+
+class _PositionalEncoding(nn.Module):
+    """
+    Positional encoding module that adds learnable positional embeddings.
+
+    Parameters
+    ----------
+    seq_length : int
+        Sequence length.
+    d_model : int
+        Dimensionality of the model.
+    """
+
+    def __init__(self, seq_length: int, d_model: int) -> None:
+        super().__init__()
+        self.seq_length = seq_length
+        self.d_model = d_model
+        self.pe = nn.Parameter(torch.zeros(1, seq_length, d_model))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.pe
+        return x
+
+
+class _Transformer(nn.Module):
+    """
+    Transformer encoder module with learnable class token and positional encoding.
+
+    Parameters
+    ----------
+    seq_length : int
+        Sequence length of the input.
+    d_model : int
+        Dimensionality of the model.
+    num_heads : int
+        Number of heads in the multihead attention.
+    feedforward_ratio : float
+        Ratio to compute the dimension of the feedforward network.
+    drop_prob : float, optional
+        Dropout probability, by default 0.5.
+    num_layers : int, optional
+        Number of transformer encoder layers, by default 4.
+    """
+
+    def __init__(
+        self,
+        seq_length: int,
+        d_model: int,
+        num_heads: int,
+        feedforward_ratio: float,
+        drop_prob: float = 0.5,
+        num_layers: int = 4,
+    ) -> None:
+        super().__init__()
+        self.cls_embedding = nn.Parameter(torch.zeros(1, 1, d_model))
+        self.pos_embedding = _PositionalEncoding(seq_length + 1, d_model)
+
+        dim_ff = int(d_model * feedforward_ratio)
+        self.dropout = nn.Dropout(drop_prob)
+        self.trans = nn.TransformerEncoder(
+            nn.TransformerEncoderLayer(
+                d_model,
+                num_heads,
+                dim_ff,
+                drop_prob,
+                batch_first=True,
+                norm_first=True,
+            ),
+            num_layers,
+            norm=nn.LayerNorm(d_model),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size = x.shape[0]
+        x = torch.cat((self.cls_embedding.expand(batch_size, -1, -1), x), dim=1)
+        x = self.pos_embedding(x)
+        x = self.dropout(x)
+        return self.trans(x)[:, 0]
+
+
+class _DenseLayers(nn.Sequential):
+    """
+    Final classification layers.
+
+    Parameters
+    ----------
+    linear_in : int
+        Input dimension to the linear layer.
+    n_classes : int
+        Number of output classes.
+    """
+
+    def __init__(self, linear_in: int, n_classes: int):
+        super().__init__(
+            nn.Flatten(),
+            nn.Linear(linear_in, n_classes),
+        )

braindecode/models/sccnet.py

@@ -0,0 +1,207 @@
+# Authors: Chun-Shu Wei
+#         Bruno Aristimunha <b.aristimunha@gmail.com> (braindecode adaptation)
+#
+# License: BSD (3-clause)
+
+import math
+from warnings import warn
+
+import torch
+from einops.layers.torch import Rearrange
+from torch import nn
+
+from braindecode.models.base import EEGModuleMixin
+from braindecode.modules import LogActivation
+
+
+class SCCNet(EEGModuleMixin, nn.Module):
+    """SCCNet from Wei, C S (2019) [sccnet]_.
+
+    Spatial component-wise convolutional network (SCCNet) for motor-imagery EEG
+    classification.
+
+    .. figure:: https://dt5vp8kor0orz.cloudfront.net/6e3ec5d729cd51fe8acc5a978db27d02a5df9e05/2-Figure1-1.png
+       :align: center
+       :alt:  Spatial component-wise convolutional network
+
+
+    1. **Spatial Component Analysis**: Performs convolution spatial filtering
+        across all EEG channels to extract spatial components, effectively
+        reducing the channel dimension.
+    2. **Spatio-Temporal Filtering**: Applies convolution across the spatial
+        components and temporal domain to capture spatio-temporal patterns.
+    3. **Temporal Smoothing (Pooling)**: Uses average pooling over time to smooth the
+       features and reduce the temporal dimension, focusing on longer-term patterns.
+    4. **Classification**: Flattens the features and applies a fully connected
+       layer.
+
+
+    Parameters
+    ----------
+    n_spatial_filters : int, optional
+        Number of spatial filters in the first convolutional layer. Default is 22.
+    n_spatial_filters_smooth : int, optional
+        Number of spatial filters used as filter in the second convolutional
+        layer. Default is 20.
+    drop_prob : float, optional
+        Dropout probability. Default is 0.5.
+    activation : nn.Module, optional
+        Activation function after the second convolutional layer. Default is
+        logarithm activation.
+
+    Notes
+    -----
+    This implementation is not guaranteed to be correct, has not been checked
+    by original authors, only reimplemented from the paper description and
+    the source that have not been tested [sccnetcode]_.
+
+
+    References
+    ----------
+    .. [sccnet] Wei, C. S., Koike-Akino, T., & Wang, Y. (2019, March). Spatial
+        component-wise convolutional network (SCCNet) for motor-imagery EEG
+        classification. In 2019 9th International IEEE/EMBS Conference on
+        Neural Engineering (NER) (pp. 328-331). IEEE.
+    .. [sccnetcode] Hsieh, C. Y., Chou, J. L., Chang, Y. H., & Wei, C. S.
+        XBrainLab: An Open-Source Software for Explainable Artificial
+        Intelligence-Based EEG Analysis. In NeurIPS 2023 AI for
+        Science Workshop.
+
+    """
+
+    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 related parameters
+        n_spatial_filters: int = 22,
+        n_spatial_filters_smooth: int = 20,
+        drop_prob: float = 0.5,
+        activation: nn.Module = LogActivation,
+        batch_norm_momentum: float = 0.1,
+    ):
+        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
+        # Parameters
+        self.n_spatial_filters = n_spatial_filters
+        self.n_spatial_filters_smooth = n_spatial_filters_smooth
+        self.drop_prob = drop_prob
+
+        # Original logical for SCCNet
+        conv_kernel_time = 0.1  # 100ms
+        pool_kernel_time = 0.5  # 500ms
+
+        # Calculate sample-based sizes from time durations
+        conv_kernel_samples = int(math.floor(self.sfreq * conv_kernel_time))
+        pool_kernel_samples = int(math.floor(self.sfreq * pool_kernel_time))
+
+        # If the input window is too short for the default kernel sizes,
+        # scale them down proportionally.
+        total_kernel_samples = conv_kernel_samples + pool_kernel_samples
+
+        if self.n_times < total_kernel_samples:
+            warning_msg = (
+                f"Input window seconds ({self.input_window_seconds:.2f}s) is smaller than the "
+                f"model's combined kernel sizes ({(total_kernel_samples / self.sfreq):.2f}s). "
+                "Scaling temporal parameters down proportionally."
+            )
+            warn(warning_msg, UserWarning, stacklevel=2)
+
+            scaling_factor = self.n_times / total_kernel_samples
+            conv_kernel_samples = int(math.floor(conv_kernel_samples * scaling_factor))
+            pool_kernel_samples = int(math.floor(pool_kernel_samples * scaling_factor))
+
+        # Ensure kernels are at least 1 sample wide
+        self.samples_100ms = max(1, conv_kernel_samples)
+        self.kernel_size_pool = max(1, pool_kernel_samples)
+
+        num_features = self._calc_num_features()
+
+        # Layers
+        self.ensure_dim = Rearrange("batch nchan times -> batch 1 nchan times")
+
+        self.activation = LogActivation() if activation is None else activation()
+
+        self.spatial_conv = nn.Conv2d(
+            in_channels=1,
+            out_channels=self.n_spatial_filters,
+            kernel_size=(self.n_chans, 1),
+        )
+
+        self.spatial_batch_norm = nn.BatchNorm2d(
+            self.n_spatial_filters, momentum=batch_norm_momentum
+        )
+
+        self.permute = Rearrange(
+            "batch filspat nchans time -> batch nchans filspat time"
+        )
+
+        self.spatial_filt_conv = nn.Conv2d(
+            in_channels=1,
+            out_channels=self.n_spatial_filters_smooth,
+            kernel_size=(self.n_spatial_filters, self.samples_100ms),
+            bias=False,
+        )
+        self.batch_norm = nn.BatchNorm2d(
+            self.n_spatial_filters_smooth, momentum=batch_norm_momentum
+        )
+
+        self.dropout = nn.Dropout(self.drop_prob)
+        self.temporal_smoothing = nn.AvgPool2d(
+            kernel_size=(1, self.kernel_size_pool),
+            stride=(1, self.samples_100ms),
+        )
+
+        self.final_layer = nn.Linear(num_features, self.n_outputs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Shape: (batch_size, n_chans, n_times)
+        x = self.ensure_dim(x)
+        # Shape: (batch_size, 1, n_chans, n_times)
+        x = self.spatial_conv(x)
+        # Shape: (batch_size, n_filters, 1, n_times)
+        x = self.spatial_batch_norm(x)
+        # Shape: (batch_size, n_filters, 1, n_times)
+        x = self.permute(x)
+        # Shape: (batch_size, 1, n_filters, n_times)
+        x = self.spatial_filt_conv(x)
+        # Shape: (batch_size, n_filters_filt, 1, n_times_reduced)
+        x = self.batch_norm(x)
+        # Shape: (batch_size, n_filters_filt, 1, n_times_reduced)
+        x = torch.pow(x, 2)
+        # Shape: (batch_size, n_filters_filt, 1, n_times_reduced)
+        x = self.dropout(x)
+        # Shape: (batch_size, n_filters_filt, 1, n_times_reduced)
+        x = self.temporal_smoothing(x)
+        # Shape: (batch_size, n_filters_filt, 1, n_times_reduced_avg_pool)
+        x = self.activation(x)
+        # Shape: (batch_size, n_filters_filt, 1, n_times_reduced_avg_pool)
+        x = x.view(x.size(0), -1)
+        # Shape: (batch_size, n_filters_filt*n_times_reduced_avg_pool)
+        x = self.final_layer(x)
+        # Shape: (batch_size, n_outputs)
+        return x
+
+    def _calc_num_features(self) -> int:
+        # Compute the number of features for the final linear layer
+        w_out_conv2 = (
+            self.n_times - self.samples_100ms + 1  # After second conv layer
+        )
+        w_out_pool = (
+            (w_out_conv2 - self.kernel_size_pool) // self.samples_100ms + 1
+            # After pooling layer
+        )
+        num_features = self.n_spatial_filters_smooth * w_out_pool
+        return num_features