Braindecode 1.8.1.dev169814997__tar.gz → 1.8.1.dev169869285__tar.gz

{braindecode-1.8.1.dev169814997/braindecode.egg-info → braindecode-1.8.1.dev169869285}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.8.1.dev169814997
+Version: 1.8.1.dev169869285
 Summary: Deep learning software to decode EEG, ECG or MEG signals
 Author-email: Robin Tibor Schirrmeister <robintibor@gmail.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Alexandre Gramfort <agramfort@meta.com>
 Maintainer-email: Alexandre Gramfort <agramfort@meta.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Robin Tibor Schirrmeister <robintibor@gmail.com>
@@ -50,7 +50,6 @@ Provides-Extra: tests
 Requires-Dist: pytest<9.1; extra == "tests"
 Requires-Dist: pytest-cov; extra == "tests"
 Requires-Dist: pytest-xdist>=3.8.0; extra == "tests"
-Requires-Dist: pytest-timeout; extra == "tests"
 Requires-Dist: codecov; extra == "tests"
 Requires-Dist: pytest_cases; extra == "tests"
 Requires-Dist: mypy; extra == "tests"

{braindecode-1.8.1.dev169814997 → braindecode-1.8.1.dev169869285}/braindecode/models/biot.py

@@ -4,10 +4,9 @@ from warnings import warn
 import torch
 import torch.nn as nn
 
-from braindecode.functional import spectral_input
+from braindecode.functional import sinusoidal_positional_encoding, spectral_input
 from braindecode.models.base import EEGModuleMixin
 from braindecode.models.util import warn_if_sfreq_differs
-from braindecode.modules import SinusoidalPositionalEncoding
 from braindecode.modules.linear_attention import LinearAttentionTransformer
 
 # -----------------------------------------------------------------------------
@@ -356,6 +355,60 @@ class _ClassificationHead(nn.Sequential):
         return out
 
 
+class _PositionalEncoding(nn.Module):
+    r"""
+    Positional Encoding.
+
+    We first create a `pe` zero matrix of shape (max_len, d_model) where max_len is the
+    the maximum length of the sequence, and emb_size is the size of the embedding.
+
+    Then we create a `position` tensor of shape (max_len, 1) with indices from 0 to max_len
+    and a `div_term` tensor of shape (d_model // 2) with the exponential of the
+    multiplication of the indices from 0 to d_model by -log(10000.0) / d_model.
+
+    For more details about the positional encoding, see the `Attention is All You Need` paper.
+
+    Parameters
+    ----------
+    emb_size: int
+        The size of the embedding layer
+    dropout: float
+        The dropout rate
+    max_len: int
+        The maximum length of the sequence
+    drop_prob : float, default=0.5
+        The dropout rate for regularization. Values should be between 0 and 1.
+
+    Returns
+    -------
+    out: Tensor
+        (batch, max_len, d_model)
+    """
+
+    def __init__(self, emb_size: int, drop_prob: float = 0.1, max_len: int = 1000):
+        super(_PositionalEncoding, self).__init__()
+
+        # Compute the positional encodings once (shared sinusoidal primitive).
+        pe = sinusoidal_positional_encoding(max_len, emb_size).unsqueeze(0)
+        self.register_buffer("pe", pe)
+        self.dropout = nn.Dropout(p=drop_prob)
+
+    def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
+        """
+        Parameters
+        ----------
+        x: FloatTensor
+            `embeddings,` shape (batch, max_len, d_model)
+
+        Returns
+        -------
+        Tensor
+            `encoder input`, shape (batch, max_len, d_model)
+        """
+        x = x + self.pe[:, : x.size(1)]
+        return self.dropout(x)
+
+
 class _BIOTEncoder(nn.Module):
     r"""
     BIOT Encoder.
@@ -416,10 +469,7 @@ class _BIOTEncoder(nn.Module):
             depth=n_layers,
             attn_layer_dropout=attn_layer_dropout,
         )
-        # persistent: released checkpoints carry the table
-        self.positional_encoding = SinusoidalPositionalEncoding(
-            1000, emb_size, drop_prob, persistent=True
-        )
+        self.positional_encoding = _PositionalEncoding(emb_size, drop_prob=drop_prob)
 
         # channel token, N_channels >= your actual channels
         self.channel_tokens = nn.Embedding(

{braindecode-1.8.1.dev169814997 → braindecode-1.8.1.dev169869285}/braindecode/models/ctnet.py

@@ -10,7 +10,6 @@ classification from Wei Zhao et al. (2024).
 from __future__ import annotations
 
 import math
-from collections import OrderedDict
 from typing import Optional
 
 import torch
@@ -21,9 +20,8 @@ from torch import Tensor, nn
 from braindecode.models.base import EEGModuleMixin
 from braindecode.modules import (
     FeedForwardBlock,
-    LearnedPositionalEncoding,
     MultiHeadAttention,
-    PostNormTransformerBlock,
+    PostNormResidual,
 )
 
 
@@ -179,47 +177,19 @@ class CTNet(EEGModuleMixin, nn.Module, license="mit"):
             activation=self.activation_patch,
         )
 
-        length = 100
-        needed = max(
-            int(self.n_times / (self.pool_size_1 * self.pool_size_1)),
-            self.sequence_length or 0,
-        )
-        if needed > length:
-            warn(
-                "the temporal dimensional is too long for this default length. "
-                "The length parameter will be automatically adjusted to "
-                "avoid inference issues."
-            )
-            length = needed
-        self.position = LearnedPositionalEncoding(
-            length, self.embed_dim, self.att_positional_drop_prob
+        self.position = _PositionalEncoding(
+            emb_size=self.embed_dim,
+            drop_prob=self.att_positional_drop_prob,
+            n_times=self.n_times,
+            pool_size=self.pool_size_1,
+            sequence_length=self.sequence_length,
         )
 
-        # ``layers`` and the ``Sequential`` around the attention keep the
-        # parameter names of saved checkpoints.
-        self.trans = nn.Sequential(
-            OrderedDict(
-                layers=nn.Sequential(
-                    *[
-                        PostNormTransformerBlock(
-                            self.embed_dim,
-                            nn.Sequential(
-                                MultiHeadAttention(self.embed_dim, self.num_heads, 0.5)
-                            ),
-                            nn.Sequential(
-                                FeedForwardBlock(
-                                    self.embed_dim,
-                                    expansion=4,
-                                    drop_p=0.5,
-                                    activation=self.activation_transformer,
-                                )
-                            ),
-                            0.5,
-                        )
-                        for _ in range(self.num_layers)
-                    ]
-                )
-            )
+        self.trans = _TransformerEncoder(
+            self.num_heads,
+            self.num_layers,
+            self.embed_dim,
+            activation=self.activation_transformer,
         )
 
         self.flatten_drop_layer = nn.Sequential(
@@ -421,3 +391,141 @@ class _PatchEmbeddingEEGNet(nn.Module):
                 x = layer(x)
         x = self.projection(x)
         return x
+
+
+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: type[nn.Module] = nn.GELU,
+    ):
+        super().__init__()
+        self.attention = PostNormResidual(
+            nn.Sequential(
+                MultiHeadAttention(dim_feedforward, num_heads, drop_prob),
+            ),
+            dim_feedforward,
+            drop_prob,
+        )
+        self.feed_forward = PostNormResidual(
+            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: type[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,
+        sequence_length: int | None = None,
+    ):
+        super().__init__()
+        self.pool_size = pool_size
+        self.n_times = n_times
+
+        # n_times / pool_size**2 is the sequence length only when both pooling
+        # sizes are equal; also cover the actual length (e.g. pool_size_2 <
+        # pool_size_1), so that the encoding is never shorter than the input
+        needed = max(int(n_times / (pool_size * pool_size)), sequence_length or 0)
+        if needed > length:
+            warn(
+                "the temporal dimensional is too long for this default length. "
+                "The length parameter will be automatically adjusted to "
+                "avoid inference issues."
+            )
+            length = needed
+
+        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-1.8.1.dev169814997 → braindecode-1.8.1.dev169869285}/braindecode/models/eegconformer.py

@@ -9,11 +9,7 @@ from einops.layers.torch import Rearrange
 from torch import Tensor, nn
 
 from braindecode.models.base import EEGModuleMixin, _too_short_input_error
-from braindecode.modules import (
-    FeedForwardBlock,
-    MultiHeadAttention,
-    PreNormTransformerBlock,
-)
+from braindecode.modules import FeedForwardBlock, MultiHeadAttention, Residual
 
 
 class EEGConformer(EEGModuleMixin, nn.Module):
@@ -254,21 +250,12 @@ class EEGConformer(EEGModuleMixin, nn.Module):
         else:
             self.final_fc_length = final_fc_length
 
-        self.transformer = nn.Sequential(
-            *[
-                PreNormTransformerBlock(
-                    n_filters_time,
-                    MultiHeadAttention(n_filters_time, num_heads, att_drop_prob),
-                    FeedForwardBlock(
-                        n_filters_time,
-                        expansion=4,
-                        drop_p=att_drop_prob,
-                        activation=activation_transfor,
-                    ),
-                    att_drop_prob,
-                )
-                for _ in range(num_layers)
-            ]
+        self.transformer = _TransformerEncoder(
+            num_layers=num_layers,
+            emb_size=n_filters_time,
+            num_heads=num_heads,
+            att_drop=att_drop_prob,
+            activation=activation_transfor,
         )
 
         self.fc = _FullyConnected(
@@ -377,6 +364,74 @@ class _PatchEmbedding(nn.Module):
         return x
 
 
+class _TransformerEncoderBlock(nn.Sequential):
+    def __init__(
+        self,
+        emb_size,
+        num_heads,
+        att_drop,
+        forward_expansion=4,
+        activation: type[nn.Module] = nn.GELU,
+    ):
+        super().__init__(
+            Residual(
+                nn.Sequential(
+                    nn.LayerNorm(emb_size),
+                    MultiHeadAttention(emb_size, num_heads, att_drop),
+                    nn.Dropout(att_drop),
+                )
+            ),
+            Residual(
+                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):
+    r"""Transformer encoder module for the transformer encoder.
+
+    Similar to the layers used in ViT.
+
+    Parameters
+    ----------
+    num_layers : int
+        Number of transformer encoder blocks.
+    emb_size : int
+        Embedding size of the transformer encoder.
+    num_heads : int
+        Number of attention heads.
+    att_drop : float
+        Dropout probability for the attention layers.
+
+    """
+
+    def __init__(
+        self,
+        num_layers,
+        emb_size,
+        num_heads,
+        att_drop,
+        activation: type[nn.Module] = nn.GELU,
+    ):
+        super().__init__(
+            *[
+                _TransformerEncoderBlock(
+                    emb_size, num_heads, att_drop, activation=activation
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+
 class _FullyConnected(nn.Module):
     def __init__(
         self,

{braindecode-1.8.1.dev169814997 → braindecode-1.8.1.dev169869285}/braindecode/models/eegdino.py

@@ -13,7 +13,7 @@ from torch import nn
 from braindecode.functional import spectral_input
 from braindecode.models.base import EEGModuleMixin
 from braindecode.models.util import warn_if_sfreq_differs
-from braindecode.modules import PatchTokenizer, PreNormBlock, TimmAttention
+from braindecode.modules import DropPath, PatchTokenizer, TimmAttention
 
 
 class EEGDINO(EEGModuleMixin, nn.Module):
@@ -236,21 +236,8 @@ class EEGDINO(EEGModuleMixin, nn.Module):
         self.emb_dim = self.patch_embedding.emb_dim
 
         self.encoder_layers = nn.ModuleList(
-            PreNormBlock(
-                self.emb_dim,
-                # BEiT attention with q/v biases, in the reference op order
-                TimmAttention(
-                    self.emb_dim,
-                    nhead,
-                    qv_bias=True,
-                    attn_drop=drop_prob,
-                    proj_drop=drop_prob,
-                    exact=True,
-                ),
-                _MLP(
-                    self.emb_dim, dim_feedforward, activation=activation, drop=drop_prob
-                ),
-                drop_path=drop_prob,
+            _TransformerEncoderLayer(
+                self.emb_dim, nhead, dim_feedforward, activation, drop_prob
             )
             for _ in range(n_layer)
         )
@@ -456,6 +443,33 @@ class _PatchEmbedding(nn.Module):
         return patch_emb + temporal_emb.permute(0, 2, 3, 1)
 
 
+class _TransformerEncoderLayer(nn.Module):
+    """Pre-norm transformer block (no LayerScale, matching the released config)."""
+
+    def __init__(
+        self, emb_dim, nhead, dim_feedforward, activation=nn.GELU, drop_prob=0.1
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(emb_dim)
+        # BEiT attention with q/v biases, in the reference op order
+        self.attn = TimmAttention(
+            emb_dim,
+            nhead,
+            qv_bias=True,
+            attn_drop=drop_prob,
+            proj_drop=drop_prob,
+            exact=True,
+        )
+        self.drop_path = DropPath(drop_prob) if drop_prob > 0.0 else nn.Identity()
+        self.norm2 = nn.LayerNorm(emb_dim)
+        self.mlp = _MLP(emb_dim, dim_feedforward, activation=activation, drop=drop_prob)
+
+    def forward(self, x):
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
 class _MLP(nn.Module):
     """Feed-forward block; ``fc1``/``fc2`` names match the released checkpoint."""
 

{braindecode-1.8.1.dev169814997 → braindecode-1.8.1.dev169869285}/braindecode/models/eegpt.py

@@ -12,7 +12,7 @@ from torch import nn
 
 from braindecode.functional import init_weights_
 from braindecode.models.base import EEGModuleMixin
-from braindecode.modules import PreNormBlock, TimmAttention, TimmMlp
+from braindecode.modules import DropPath, TimmAttention, TimmMlp
 from braindecode.modules.convolution import Conv1dWithConstraint
 from braindecode.modules.linear import LinearWithConstraint
 from braindecode.util import resolve_montage_name
@@ -734,6 +734,84 @@ def _apply_mask_t(mask_t, x):
     return masked_x
 
 
+class _Block(nn.Module):
+    """
+    Transformer Block consisting of Self-Attention and MLP.
+
+    This is a standard Transformer encoder block that applies:
+    1. Layer Normalization
+    2. Multi-Head Self-Attention
+    3. Residual Connection with DropPath
+    4. Layer Normalization
+    5. MLP
+    6. Residual Connection with DropPath
+
+    Parameters
+    ----------
+    dim : int
+        Input and output embedding dimension.
+    num_heads : int
+        Number of attention heads.
+    mlp_ratio : float, default=4.0
+        Ratio of MLP hidden dimension to embedding dimension.
+    qkv_bias : bool, default=False
+        If True, add a learnable bias to query, key, value projections.
+    drop : float, default=0.0
+        Dropout probability for MLP and projection layers.
+    attn_drop : float, default=0.0
+        Dropout probability for attention weights.
+    drop_path : float, default=0.0
+        Stochastic depth rate.
+    act_layer : nn.Module, default=nn.GELU
+        Activation function.
+    norm_layer : nn.Module, default=nn.LayerNorm
+        Normalization layer.
+    return_attention : bool, default=False
+        If True, returns attention weights instead of the output tensor.
+    """
+
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+        return_attention=False,
+    ):
+        super().__init__()
+        self.return_attention = return_attention
+        self.norm1 = norm_layer(dim)
+        self.attn = TimmAttention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = TimmMlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+
+    def forward(self, x):
+        y = self.attn(self.norm1(x), need_weights=self.return_attention)
+        if self.return_attention:
+            return y
+        x = x + self.drop_path(y)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
 class _PatchEmbed(nn.Module):
     """
     Splits the EEG signal into patches and embeds them.
@@ -930,7 +1008,6 @@ class _EEGTransformer(nn.Module):
         self.embed_dim = embed_dim
         self.embed_num = embed_num
         self.num_heads = num_heads
-        self.return_attention_layer = return_attention_layer
 
         self.patch_embed = patch_module(
             n_chans=n_chans,
@@ -952,22 +1029,16 @@ class _EEGTransformer(nn.Module):
 
         self.blocks = nn.ModuleList(
             [
-                PreNormBlock(
-                    embed_dim,
-                    TimmAttention(
-                        embed_dim,
-                        num_heads=num_heads,
-                        qkv_bias=qkv_bias,
-                        attn_drop=attn_drop_rate,
-                        proj_drop=drop_rate,
-                    ),
-                    TimmMlp(
-                        in_features=embed_dim,
-                        hidden_features=int(embed_dim * mlp_ratio),
-                        drop=drop_rate,
-                    ),
-                    norm_layer=norm_layer,
+                _Block(
+                    dim=embed_dim,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    drop=drop_rate,
+                    attn_drop=attn_drop_rate,
                     drop_path=dpr[i],
+                    norm_layer=norm_layer,
+                    return_attention=(i + 1) == return_attention_layer,
                 )
                 for i in range(depth)
             ]
@@ -1092,9 +1163,9 @@ class _EEGTransformer(nn.Module):
 
         # -- fwd prop
         for i, blk in enumerate(self.blocks):
-            if i + 1 == self.return_attention_layer:
-                return blk.attn(blk.norm1(x), need_weights=True)
             x = blk(x)
+            if blk.return_attention:
+                return x
 
         # Extract only the summary tokens (used for global representation)
         # x shape: (batch * n_patches, embed_num, embed_dim)