braindecode 1.5.0.dev179040812__tar.gz → 1.5.0.dev182195895__tar.gz

{braindecode-1.5.0.dev179040812/braindecode.egg-info → braindecode-1.5.0.dev182195895}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.5.0.dev179040812
+Version: 1.5.0.dev182195895
 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>

{braindecode-1.5.0.dev179040812 → braindecode-1.5.0.dev182195895}/braindecode/models/attentionbasenet.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+# Authors: Sarthak Tayal <sarthaktayal2@gmail.com>
+#
+# License: BSD (3-clause)
 import math
 
 from einops.layers.torch import Rearrange
@@ -275,8 +278,6 @@ class AttentionBaseNet(EEGModuleMixin, nn.Module):
         activation: type[nn.Module] = nn.ELU,
         extra_params: bool = False,
     ):
-        super(AttentionBaseNet, self).__init__()
-
         super().__init__(
             n_outputs=n_outputs,
             n_chans=n_chans,

{braindecode-1.5.0.dev179040812 → braindecode-1.5.0.dev182195895}/braindecode/models/biot.py

@@ -439,7 +439,7 @@ class _BIOTEncoder(nn.Module):
         self.channel_tokens = nn.Embedding(
             num_embeddings=n_chans, embedding_dim=emb_size
         )
-        self.index = nn.Parameter(torch.LongTensor(range(n_chans)), requires_grad=False)
+        self.register_buffer("index", torch.arange(n_chans, dtype=torch.long))
 
     def stft(self, sample):
         """

{braindecode-1.5.0.dev179040812 → braindecode-1.5.0.dev182195895}/braindecode/models/syncnet.py

@@ -1,3 +1,7 @@
+# Authors: Sarthak Tayal <sarthaktayal2@gmail.com>
+#
+# License: BSD (3-clause)
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -51,11 +55,11 @@ class SyncNet(EEGModuleMixin, nn.Module):
         The initialization range for omega parameters using uniform
         distribution. Default is (0, 1).
     beta_init_values : tuple of float, optional
-        The initialization range for beta parameters using uniform
-        distribution. Default is (0, 1). Default is (0, 0.05).
+        The initialization range for beta (decay) parameters using uniform
+        distribution. Default is (0, 0.05).
     phase_init_values : tuple of float, optional
-        The initialization range for phase parameters using `normal`
-        distribution. Default is (0, 1). Default is (0, 0.05).
+        The initialization mean and standard deviation for phase
+        parameters using normal distribution. Default is (0, 0.05).
 
 
     Notes
@@ -146,12 +150,12 @@ class SyncNet(EEGModuleMixin, nn.Module):
         # Phase Shift
         self.phi_ini = nn.Parameter(
             torch.FloatTensor(1, 1, self.n_chans, self.num_filters).normal_(
-                self.beta_init_values[0], self.beta_init_values[1]
+                self.phase_init_values[0], self.phase_init_values[1]
             )
         )
         self.beta = nn.Parameter(
             torch.FloatTensor(1, 1, 1, self.num_filters).uniform_(
-                self.phase_init_values[0], self.phase_init_values[1]
+                self.beta_init_values[0], self.beta_init_values[1]
             )
         )
 
@@ -185,21 +189,21 @@ class SyncNet(EEGModuleMixin, nn.Module):
         # Output: (batch_size, n_chans, 1, n_times)
 
         # Compute the oscillatory component
+        # Shape: (1, filter_width, n_chans, num_filters)
         W_osc = self.amplitude * torch.cos(self.t * self.omega + self.phi_ini)
-        # W_osc is (1, filter_width, n_chans, 1)
 
         # Compute the decay component
-        t_squared = torch.pow(self.t, 2)  # Shape: (filter_width,)
-        t_squared_beta = t_squared * self.beta  # Shape: (filter_width, num_filters)
+        # Shape: (1, filter_width, 1, num_filters)
+        t_squared = torch.pow(self.t, 2)
+        t_squared_beta = t_squared * self.beta
         W_decay = torch.exp(-t_squared_beta)
-        # W_osc is (1, filter_width, 1, 1)
 
         # Combine oscillatory and decay components
-        # W shape: (1, n_chans, num_filters, filter_width)
+        # Shape: (1, filter_width, n_chans, num_filters)
         W = W_osc * W_decay
-        # W shape will be: (1, filter_width, n_chans, 1)
 
-        W = W.view(self.num_filters, self.n_chans, 1, self.filter_width)
+        # Permute to conv2d weight shape (out_channels, in_channels, kH, kW)
+        W = W.permute(3, 2, 0, 1).contiguous()
 
         # Apply convolution
         x_padded = self.pad_input(x.float())

{braindecode-1.5.0.dev179040812 → braindecode-1.5.0.dev182195895}/braindecode/models/tsinception.py

@@ -1,4 +1,4 @@
-# Authors: Bruno Aristimunha <b.aristimunha>
+# Authors: Bruno Aristimunha <b.aristimunha@gmail.com>
 #
 # License: BSD (3-clause)
 

{braindecode-1.5.0.dev179040812 → braindecode-1.5.0.dev182195895}/braindecode/modules/attention.py

@@ -837,6 +837,18 @@ class CATLite(nn.Module):
 class MultiHeadAttention(nn.Module):
     """Multi-head self-attention block.
 
+    Uses ``F.scaled_dot_product_attention`` for optimized attention
+    kernels (flash-attention on CUDA, memory-efficient on other devices).
+
+    Parameters
+    ----------
+    emb_size : int
+        The embedding dimension.
+    num_heads : int
+        Number of attention heads. Must evenly divide ``emb_size``.
+    dropout : float, optional
+        Dropout probability applied to attention weights. Default: 0.0.
+
     Examples
     --------
     >>> import torch
@@ -848,40 +860,57 @@ class MultiHeadAttention(nn.Module):
     torch.Size([2, 10, 32])
     """
 
-    def __init__(self, emb_size, num_heads, dropout):
+    def __init__(self, emb_size, num_heads, dropout=0.0):
         super().__init__()
+        if emb_size % num_heads != 0:
+            raise ValueError(
+                f"emb_size ({emb_size}) must be divisible by num_heads ({num_heads})."
+            )
         self.emb_size = emb_size
         self.num_heads = num_heads
+        self.head_dim = emb_size // num_heads
         self.keys = nn.Linear(emb_size, emb_size)
         self.queries = nn.Linear(emb_size, emb_size)
         self.values = nn.Linear(emb_size, emb_size)
-        self.att_drop = nn.Dropout(dropout)
+        self.att_drop = dropout
         self.projection = nn.Linear(emb_size, emb_size)
 
         self.rearrange_stack = Rearrange(
-            "b n (h d) -> b h n d",
-            h=num_heads,
+            "batch seq (heads head_dim) -> batch heads seq head_dim",
+            heads=num_heads,
         )
         self.rearrange_unstack = Rearrange(
-            "b h n d -> b n (h d)",
+            "batch heads seq head_dim -> batch seq (heads head_dim)",
         )
 
     def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
+        """Forward pass.
+
+        Parameters
+        ----------
+        x : Tensor
+            Input tensor of shape ``(batch, seq, emb_size)``.
+        mask : Tensor, optional
+            Attention mask following PyTorch SDPA convention: for boolean
+            masks ``True`` means **ignore** that position; for float
+            masks the values are **added** to attention scores before
+            softmax.
+        """
         queries = self.rearrange_stack(self.queries(x))
         keys = self.rearrange_stack(self.keys(x))
         values = self.rearrange_stack(self.values(x))
-        energy = torch.einsum("bhqd, bhkd -> bhqk", queries, keys)
-        if mask is not None:
-            fill_value = float("-inf")
-            energy = energy.masked_fill(~mask, fill_value)
-
-        scaling = self.emb_size ** (1 / 2)
-        att = F.softmax(energy / scaling, dim=-1)
-        att = self.att_drop(att)
-        out = torch.einsum("bhal, bhlv -> bhav ", att, values)
+
+        dp = self.att_drop if self.training else 0.0
+        out = F.scaled_dot_product_attention(
+            queries,
+            keys,
+            values,
+            attn_mask=mask,
+            dropout_p=dp,
+        )
+
         out = self.rearrange_unstack(out)
-        out = self.projection(out)
-        return out
+        return self.projection(out)
 
 
 class CrissCrossTransformerEncoderLayer(nn.Module):

braindecode-1.5.0.dev182195895/braindecode/version.py

@@ -0,0 +1 @@
+__version__ = "1.5.0.dev182195895"

{braindecode-1.5.0.dev179040812 → braindecode-1.5.0.dev182195895/braindecode.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.5.0.dev179040812
+Version: 1.5.0.dev182195895
 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>

{braindecode-1.5.0.dev179040812 → braindecode-1.5.0.dev182195895}/docs/whats_new.rst

@@ -28,10 +28,19 @@ Current 1.5.0 (GitHub)
 Enhancements
 ============
 
+- Use ``F.scaled_dot_product_attention`` in :class:`braindecode.modules.MultiHeadAttention`,
+  enabling optimized attention kernels (flash-attention on CUDA,
+  memory-efficient backends on other devices).
+  By `Léo Burgund`_ and `Bruno Aristimunha`_.
+  (:gh:`902`)
+
 API and behavior changes
 ========================
 
-- None yet
+- :class:`braindecode.modules.MultiHeadAttention` now follows PyTorch's SDPA mask
+  convention: boolean masks use ``True`` to **ignore** a position (previously
+  ``True`` meant keep). The scaling factor is now ``1/sqrt(head_dim)`` instead of
+  ``1/sqrt(emb_size)``. (:gh:`902`)
 
 Requirements
 ============
@@ -41,6 +50,16 @@ Requirements
 Bug fixes
 ==========
 
+- Fix :class:`braindecode.models.SyncNet` swapped parameter initialization where
+  ``phi_ini`` (phase shift) was using ``beta_init_values`` and ``beta`` (decay) was
+  using ``phase_init_values``, replaced incorrect ``.view()`` reshape with ``.permute()``
+  for proper conv2d filter weight layout, and fixed duplicate default values in docstring
+  (by `Sarthak Tayal`_)
+- Fix :class:`braindecode.models.AttentionBaseNet` redundant
+  ``super().__init__()`` call that ran the parent ``nn.Module.__init__`` twice
+  (by `Sarthak Tayal`_)
+- Fix incomplete author email in :class:`braindecode.models.TSception` header
+  (by `Sarthak Tayal`_)
 - Fix a time-of-check-time-of-use race in
   :func:`braindecode.datasets.base._zarr_to_memmap` that caused
   concurrent workers to repeatedly ``rename``-replace the published
@@ -50,6 +69,9 @@ Bug fixes
   is never replaced, making the cache safe under arbitrary
   concurrent access on local POSIX, NFSv3, Lustre and SMB
   (:gh:`986` by `Pierre Guetschel`_)
+- Register :class:`braindecode.models.BIOT` encoder ``index`` as a non-trainable
+  buffer instead of a parameter (``torch.long``), so it is treated as module
+  state rather than trainable weights (:gh:`988` by `Pierre Guetschel`_)
 
 Code health
 ============

braindecode-1.5.0.dev179040812/braindecode/version.py

@@ -1 +0,0 @@
-__version__ = "1.5.0.dev179040812"