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

braindecode/models/ctnet.py

@@ -0,0 +1,450 @@
+"""
+CTNet: a convolutional transformer network for EEG-based motor imagery
+classification from Wei Zhao et al. (2024).
+"""
+
+# Authors: Wei Zhao <zhaowei701@163.com>
+#          Bruno Aristimunha <b.aristimunha@gmail.com> (braindecode adaptation)
+# License: MIT
+
+from __future__ import annotations
+
+import math
+
+import torch
+from einops.layers.torch import Rearrange
+from mne.utils import warn
+from torch import Tensor, nn
+
+from braindecode.models.base import EEGModuleMixin
+from braindecode.modules import (
+    FeedForwardBlock,
+    MultiHeadAttention,
+)
+
+
+class CTNet(EEGModuleMixin, nn.Module):
+    """CTNet from Zhao, W et al (2024) [ctnet]_.
+
+     A Convolutional Transformer Network for EEG-Based Motor Imagery Classification
+
+     .. figure:: https://raw.githubusercontent.com/snailpt/CTNet/main/architecture.png
+        :align: center
+        :alt: CTNet Architecture
+
+    CTNet is an end-to-end neural network architecture designed for classifying motor imagery (MI) tasks from EEG signals.
+    The model combines convolutional neural networks (CNNs) with a Transformer encoder to capture both local and global temporal dependencies in the EEG data.
+
+    The architecture consists of three main components:
+
+    1. **Convolutional Module**:
+        - Apply EEGNetV4 to perform some feature extraction, denoted here as
+        _PatchEmbeddingEEGNet module.
+
+    2. **Transformer Encoder Module**:
+        - Utilizes multi-head self-attention mechanisms as EEGConformer but
+        with residual blocks.
+
+    3. **Classifier Module**:
+        - Combines features from both the convolutional module
+        and the Transformer encoder.
+        - Flattens the combined features and applies dropout for regularization.
+        - Uses a fully connected layer to produce the final classification output.
+
+    Parameters
+    ----------
+    activation : nn.Module, default=nn.GELU
+        Activation function to use in the network.
+    heads : int, default=4
+        Number of attention heads in the Transformer encoder.
+    emb_size : int, default=40
+        Embedding size (dimensionality) for the Transformer encoder.
+    depth : int, default=6
+        Number of encoder layers in the Transformer.
+    n_filters_time : int, default=20
+        Number of temporal filters in the first convolutional layer.
+    kernel_size : int, default=64
+        Kernel size for the temporal convolutional layer.
+    depth_multiplier : int, default=2
+        Multiplier for the number of depth-wise convolutional filters.
+    pool_size_1 : int, default=8
+        Pooling size for the first average pooling layer.
+    pool_size_2 : int, default=8
+        Pooling size for the second average pooling layer.
+    drop_prob_cnn : float, default=0.3
+        Dropout probability after convolutional layers.
+    drop_prob_posi : float, default=0.1
+        Dropout probability for the positional encoding in the Transformer.
+    drop_prob_final : float, default=0.5
+        Dropout probability before the final classification layer.
+
+    Notes
+    -----
+    This implementation is adapted from the original CTNet source code
+    [ctnetcode]_ to comply with Braindecode's model standards.
+
+    References
+    ----------
+    .. [ctnet] Zhao, W., Jiang, X., Zhang, B., Xiao, S., & Weng, S. (2024).
+        CTNet: a convolutional transformer network for EEG-based motor imagery
+        classification. Scientific Reports, 14(1), 20237.
+    .. [ctnetcode] Zhao, W., Jiang, X., Zhang, B., Xiao, S., & Weng, S. (2024).
+        CTNet source code:
+        https://github.com/snailpt/CTNet
+    """
+
+    def __init__(
+        self,
+        # Base arguments
+        n_outputs=None,
+        n_chans=None,
+        sfreq=None,
+        chs_info=None,
+        n_times=None,
+        input_window_seconds=None,
+        # Model specific arguments
+        activation_patch: nn.Module = nn.ELU,
+        activation_transformer: nn.Module = nn.GELU,
+        drop_prob_cnn: float = 0.3,
+        drop_prob_posi: float = 0.1,
+        drop_prob_final: float = 0.5,
+        # other parameters
+        heads: int = 4,
+        emb_size: int = 40,
+        depth: int = 6,
+        n_filters_time: int = 20,
+        kernel_size: int = 64,
+        depth_multiplier: int = 2,
+        pool_size_1: int = 8,
+        pool_size_2: int = 8,
+    ):
+        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.emb_size = emb_size
+        self.activation_patch = activation_patch
+        self.activation_transformer = activation_transformer
+
+        self.n_filters_time = n_filters_time
+        self.drop_prob_cnn = drop_prob_cnn
+        self.pool_size_1 = pool_size_1
+        self.pool_size_2 = pool_size_2
+        self.depth_multiplier = depth_multiplier
+        self.kernel_size = kernel_size
+        self.drop_prob_posi = drop_prob_posi
+        self.drop_prob_final = drop_prob_final
+
+        # n_times - pool_size_1 / p
+        sequence_length = math.floor(
+            (
+                math.floor((self.n_times - self.pool_size_1) / self.pool_size_1 + 1)
+                - self.pool_size_2
+            )
+            / self.pool_size_2
+            + 1
+        )
+
+        # Layers
+        self.ensuredim = Rearrange("batch nchans time -> batch 1 nchans time")
+        self.flatten = nn.Flatten()
+
+        self.cnn = _PatchEmbeddingEEGNet(
+            n_filters_time=self.n_filters_time,
+            kernel_size=self.kernel_size,
+            depth_multiplier=self.depth_multiplier,
+            pool_size_1=self.pool_size_1,
+            pool_size_2=self.pool_size_2,
+            drop_prob=self.drop_prob_cnn,
+            n_chans=self.n_chans,
+            activation=self.activation_patch,
+        )
+
+        self.position = _PositionalEncoding(
+            emb_size=emb_size,
+            drop_prob=self.drop_prob_posi,
+            n_times=self.n_times,
+            pool_size=self.pool_size_1,
+        )
+
+        self.trans = _TransformerEncoder(
+            heads, depth, emb_size, activation=self.activation_transformer
+        )
+
+        self.flatten_drop_layer = nn.Sequential(
+            nn.Flatten(),
+            nn.Dropout(p=self.drop_prob_final),
+        )
+
+        self.final_layer = nn.Linear(
+            in_features=emb_size * sequence_length, out_features=self.n_outputs
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Forward pass of the CTNet model.
+
+        Parameters
+        ----------
+        x : Tensor
+            Input tensor of shape (batch_size, n_channels, n_times).
+
+        Returns
+        -------
+        Tensor
+            Output with shape (batch_size, n_outputs).
+        """
+        x = self.ensuredim(x)
+        cnn = self.cnn(x)
+        cnn = cnn * math.sqrt(self.emb_size)
+        cnn = self.position(cnn)
+        trans = self.trans(cnn)
+        features = cnn + trans
+        flatten_feature = self.flatten(features)
+        out = self.final_layer(flatten_feature)
+        return out
+
+
+class _PatchEmbeddingEEGNet(nn.Module):
+    def __init__(
+        self,
+        n_filters_time: int = 16,
+        kernel_size: int = 64,
+        depth_multiplier: int = 2,
+        pool_size_1: int = 8,
+        pool_size_2: int = 8,
+        drop_prob: float = 0.3,
+        n_chans: int = 22,
+        activation: nn.Module = nn.ELU,
+    ):
+        super().__init__()
+        n_filters_out = depth_multiplier * n_filters_time
+        self.eegnet_module = nn.Sequential(
+            # Temporal convolution
+            nn.Conv2d(
+                in_channels=1,
+                out_channels=n_filters_time,
+                kernel_size=(1, kernel_size),
+                stride=(1, 1),
+                padding="same",
+                bias=False,
+            ),
+            nn.BatchNorm2d(n_filters_time),
+            # Channel depth-wise convolution
+            nn.Conv2d(
+                in_channels=n_filters_time,
+                out_channels=n_filters_out,
+                kernel_size=(n_chans, 1),
+                stride=(1, 1),
+                groups=n_filters_time,
+                padding="valid",
+                bias=False,
+            ),
+            nn.BatchNorm2d(n_filters_out),
+            activation(),
+            # First average pooling
+            nn.AvgPool2d(kernel_size=(1, pool_size_1)),
+            nn.Dropout(drop_prob),
+            # Spatial convolution
+            nn.Conv2d(
+                in_channels=n_filters_out,
+                out_channels=n_filters_out,
+                kernel_size=(1, 16),
+                padding="same",
+                bias=False,
+            ),
+            nn.BatchNorm2d(n_filters_out),
+            activation(),
+            # Second average pooling
+            nn.AvgPool2d(kernel_size=(1, pool_size_2)),
+            nn.Dropout(drop_prob),
+        )
+
+        self.projection = nn.Sequential(
+            Rearrange("b e h w -> b (h w) e"),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Forward pass of the Patch Embedding CNN.
+
+        Parameters
+        ----------
+        x : Tensor
+            Input tensor of shape (batch_size, 1, n_channels, n_times).
+
+        Returns
+        -------
+        Tensor
+            Embedded patches of shape (batch_size, num_patches, embedding_dim).
+        """
+        x = self.eegnet_module(x)
+        x = self.projection(x)
+        return x
+
+
+class _ResidualAdd(nn.Module):
+    def __init__(self, module: nn.Module, emb_size: int, drop_p: float):
+        super().__init__()
+        self.module = module
+        self.drop = nn.Dropout(drop_p)
+        self.layernorm = nn.LayerNorm(emb_size)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Forward pass with residual connection.
+
+        Parameters
+        ----------
+        x : Tensor
+            Input tensor.
+        **kwargs : Any
+            Additional arguments.
+
+        Returns
+        -------
+        Tensor
+            Output tensor after applying residual connection.
+        """
+        res = self.module(x)
+        out = self.layernorm(self.drop(res) + x)
+        return out
+
+
+class _TransformerEncoderBlock(nn.Module):
+    def __init__(
+        self,
+        dim_feedforward: int,
+        num_heads: int = 4,
+        drop_prob: float = 0.5,
+        forward_expansion: int = 4,
+        forward_drop_p: float = 0.5,
+        activation: nn.Module = nn.GELU,
+    ):
+        super().__init__()
+        self.attention = _ResidualAdd(
+            nn.Sequential(
+                MultiHeadAttention(dim_feedforward, num_heads, drop_prob),
+            ),
+            dim_feedforward,
+            drop_prob,
+        )
+        self.feed_forward = _ResidualAdd(
+            nn.Sequential(
+                FeedForwardBlock(
+                    dim_feedforward,
+                    expansion=forward_expansion,
+                    drop_p=forward_drop_p,
+                    activation=activation,
+                ),
+            ),
+            dim_feedforward,
+            drop_prob,
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Forward pass of the transformer encoder block.
+
+        Parameters
+        ----------
+        x : Tensor
+            Input tensor.
+        **kwargs : Any
+            Additional arguments.
+
+        Returns
+        -------
+        Tensor
+            Output tensor after transformer encoder block.
+        """
+        x = self.attention(x)
+        x = self.feed_forward(x)
+        return x
+
+
+class _TransformerEncoder(nn.Module):
+    def __init__(
+        self,
+        nheads: int,
+        depth: int,
+        dim_feedforward: int,
+        activation: nn.Module = nn.GELU,
+    ):
+        super().__init__()
+        self.layers = nn.Sequential(
+            *[
+                _TransformerEncoderBlock(
+                    dim_feedforward=dim_feedforward,
+                    num_heads=nheads,
+                    activation=activation,
+                )
+                for _ in range(depth)
+            ]
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Forward pass of the transformer encoder.
+
+        Parameters
+        ----------
+        x : Tensor
+            Input tensor.
+
+        Returns
+        -------
+        Tensor
+            Output tensor after transformer encoder.
+        """
+        return self.layers(x)
+
+
+class _PositionalEncoding(nn.Module):
+    def __init__(
+        self,
+        n_times: int,
+        emb_size: int,
+        length: int = 100,
+        drop_prob: float = 0.1,
+        pool_size: int = 8,
+    ):
+        super().__init__()
+        self.pool_size = pool_size
+        self.n_times = n_times
+
+        if int(n_times / (pool_size * pool_size)) > length:
+            warn(
+                "the temporal dimensional is too long for this default length. "
+                "The length parameter will be automatically adjusted to "
+                "avoid inference issues."
+            )
+            length = int(n_times / (pool_size * pool_size))
+
+        self.dropout = nn.Dropout(drop_prob)
+        self.encoding = nn.Parameter(torch.randn(1, length, emb_size))
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Forward pass of the positional encoding.
+
+        Parameters
+        ----------
+        x : Tensor
+            Input tensor of shape (batch_size, sequence_length, embedding_dim).
+
+        Returns
+        -------
+        Tensor
+            Tensor with positional encoding added.
+        """
+        seq_length = x.size(1)
+        encoding = self.encoding[:, :seq_length, :]
+        x = x + encoding
+        return self.dropout(x)

braindecode/models/deep4.py

@@ -5,15 +5,22 @@
 from einops.layers.torch import Rearrange
 from torch import nn
 from torch.nn import init
-from torch.nn.functional import elu
 
-from .base import EEGModuleMixin, deprecated_args
-from .functions import identity, squeeze_final_output
-from .modules import AvgPool2dWithConv, CombinedConv, Ensure4d, Expression
+from braindecode.models.base import EEGModuleMixin
+from braindecode.modules import (
+    AvgPool2dWithConv,
+    CombinedConv,
+    Ensure4d,
+    SqueezeFinalOutput,
+)
 
 
 class Deep4Net(EEGModuleMixin, nn.Sequential):
-    """Deep ConvNet model from Schirrmeister et al 2017.
+    """Deep ConvNet model from Schirrmeister et al (2017) [Schirrmeister2017]_.
+
+     .. figure:: https://onlinelibrary.wiley.com/cms/asset/fc200ccc-d8c4-45b4-8577-56ce4d15999a/hbm23730-fig-0001-m.jpg
+        :align: center
+        :alt: CTNet Architecture
 
     Model described in [Schirrmeister2017]_.
 
@@ -44,13 +51,13 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
         Number of temporal filters in layer 4.
     filter_length_4: int
         Length of the temporal filter in layer 4.
-    first_conv_nonlin: callable
+    activation_first_conv_nonlin: nn.Module, default is nn.ELU
         Non-linear activation function to be used after convolution in layer 1.
     first_pool_mode: str
         Pooling mode in layer 1. "max" or "mean".
     first_pool_nonlin: callable
         Non-linear activation function to be used after pooling in layer 1.
-    later_conv_nonlin: callable
+    activation_later_conv_nonlin: nn.Module, default is nn.ELU
         Non-linear activation function to be used after convolution in later layers.
     later_pool_mode: str
         Pooling mode in later layers. "max" or "mean".
@@ -67,12 +74,6 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
         Momentum for BatchNorm2d.
     stride_before_pool: bool
         Stride before pooling.
-    in_chans :
-        Alias for n_chans.
-    n_classes:
-        Alias for n_outputs.
-    input_window_samples :
-        Alias for n_times.
 
 
     References
@@ -87,47 +88,37 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
     """
 
     def __init__(
-            self,
-            n_chans=None,
-            n_outputs=None,
-            n_times=None,
-            final_conv_length="auto",
-            n_filters_time=25,
-            n_filters_spat=25,
-            filter_time_length=10,
-            pool_time_length=3,
-            pool_time_stride=3,
-            n_filters_2=50,
-            filter_length_2=10,
-            n_filters_3=100,
-            filter_length_3=10,
-            n_filters_4=200,
-            filter_length_4=10,
-            first_conv_nonlin=elu,
-            first_pool_mode="max",
-            first_pool_nonlin=identity,
-            later_conv_nonlin=elu,
-            later_pool_mode="max",
-            later_pool_nonlin=identity,
-            drop_prob=0.5,
-            split_first_layer=True,
-            batch_norm=True,
-            batch_norm_alpha=0.1,
-            stride_before_pool=False,
-            chs_info=None,
-            input_window_seconds=None,
-            sfreq=None,
-            in_chans=None,
-            n_classes=None,
-            input_window_samples=None,
-            add_log_softmax=True,
+        self,
+        n_chans=None,
+        n_outputs=None,
+        n_times=None,
+        final_conv_length="auto",
+        n_filters_time=25,
+        n_filters_spat=25,
+        filter_time_length=10,
+        pool_time_length=3,
+        pool_time_stride=3,
+        n_filters_2=50,
+        filter_length_2=10,
+        n_filters_3=100,
+        filter_length_3=10,
+        n_filters_4=200,
+        filter_length_4=10,
+        activation_first_conv_nonlin: nn.Module = nn.ELU,
+        first_pool_mode="max",
+        first_pool_nonlin: nn.Module = nn.Identity,
+        activation_later_conv_nonlin: nn.Module = nn.ELU,
+        later_pool_mode="max",
+        later_pool_nonlin: nn.Module = nn.Identity,
+        drop_prob=0.5,
+        split_first_layer=True,
+        batch_norm=True,
+        batch_norm_alpha=0.1,
+        stride_before_pool=False,
+        chs_info=None,
+        input_window_seconds=None,
+        sfreq=None,
     ):
-        n_chans, n_outputs, n_times = deprecated_args(
-            self,
-            ('in_chans', 'n_chans', in_chans, n_chans),
-            ('n_classes', 'n_outputs', n_classes, n_outputs),
-            ('input_window_samples', 'n_times', input_window_samples, n_times),
-        )
         super().__init__(
             n_outputs=n_outputs,
             n_chans=n_chans,
@@ -135,10 +126,9 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
             n_times=n_times,
             input_window_seconds=input_window_seconds,
             sfreq=sfreq,
-            add_log_softmax=add_log_softmax,
         )
         del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
-        del in_chans, n_classes, input_window_samples
+
         if final_conv_length == "auto":
             assert self.n_times is not None
         self.final_conv_length = final_conv_length
@@ -153,10 +143,10 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
         self.filter_length_3 = filter_length_3
         self.n_filters_4 = n_filters_4
         self.filter_length_4 = filter_length_4
-        self.first_nonlin = first_conv_nonlin
+        self.first_nonlin = activation_first_conv_nonlin
         self.first_pool_mode = first_pool_mode
         self.first_pool_nonlin = first_pool_nonlin
-        self.later_conv_nonlin = later_conv_nonlin
+        self.later_conv_nonlin = activation_later_conv_nonlin
         self.later_pool_mode = later_pool_mode
         self.later_pool_nonlin = later_pool_nonlin
         self.drop_prob = drop_prob
@@ -174,7 +164,7 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
             "conv_time.bias": "conv_time_spat.conv_time.bias",
             "conv_spat.bias": "conv_time_spat.conv_spat.bias",
             "conv_classifier.weight": "final_layer.conv_classifier.weight",
-            "conv_classifier.bias": "final_layer.conv_classifier.bias"
+            "conv_classifier.bias": "final_layer.conv_classifier.bias",
         }
 
         if self.stride_before_pool:
@@ -223,17 +213,17 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
                     eps=1e-5,
                 ),
             )
-        self.add_module("conv_nonlin", Expression(self.first_nonlin))
+        self.add_module("conv_nonlin", self.first_nonlin())
         self.add_module(
             "pool",
             first_pool_class(
                 kernel_size=(self.pool_time_length, 1), stride=(pool_stride, 1)
             ),
         )
-        self.add_module("pool_nonlin", Expression(self.first_pool_nonlin))
+        self.add_module("pool_nonlin", self.first_pool_nonlin())
 
         def add_conv_pool_block(
-                model, n_filters_before, n_filters, filter_length, block_nr
+            model, n_filters_before, n_filters, filter_length, block_nr
         ):
             suffix = "_{:d}".format(block_nr)
             self.add_module("drop" + suffix, nn.Dropout(p=self.drop_prob))
@@ -257,7 +247,7 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
                         eps=1e-5,
                     ),
                 )
-            self.add_module("nonlin" + suffix, Expression(self.later_conv_nonlin))
+            self.add_module("nonlin" + suffix, self.later_conv_nonlin())
 
             self.add_module(
                 "pool" + suffix,
@@ -266,7 +256,7 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
                     stride=(pool_stride, 1),
                 ),
             )
-            self.add_module("pool_nonlin" + suffix, Expression(self.later_pool_nonlin))
+            self.add_module("pool_nonlin" + suffix, self.later_pool_nonlin())
 
         add_conv_pool_block(
             self, n_filters_conv, self.n_filters_2, self.filter_length_2, 2
@@ -286,17 +276,17 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
         # Incorporating classification module and subsequent ones in one final layer
         module = nn.Sequential()
 
-        module.add_module("conv_classifier",
-                          nn.Conv2d(
-                            self.n_filters_4,
-                            self.n_outputs,
-                            (self.final_conv_length, 1),
-                            bias=True, ))
-
-        if self.add_log_softmax:
-            module.add_module("logsoftmax", nn.LogSoftmax(dim=1))
+        module.add_module(
+            "conv_classifier",
+            nn.Conv2d(
+                self.n_filters_4,
+                self.n_outputs,
+                (self.final_conv_length, 1),
+                bias=True,
+            ),
+        )
 
-        module.add_module("squeeze", Expression(squeeze_final_output))
+        module.add_module("squeeze", SqueezeFinalOutput())
 
         self.add_module("final_layer", module)
 
@@ -329,5 +319,4 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
         init.xavier_uniform_(self.final_layer.conv_classifier.weight, gain=1)
         init.constant_(self.final_layer.conv_classifier.bias, 0)
 
-        # Start in eval mode
-        self.eval()
+        self.train()