braindecode 1.6.1.dev1056__tar.gz → 1.6.1.dev1065__tar.gz

{braindecode-1.6.1.dev1056/braindecode.egg-info → braindecode-1.6.1.dev1065}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.6.1.dev1056
+Version: 1.6.1.dev1065
 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.6.1.dev1056 → braindecode-1.6.1.dev1065}/braindecode/models/__init__.py

@@ -55,6 +55,7 @@ from .sleep_stager_chambon_2018 import SleepStagerChambon2018
 from .sparcnet import SPARCNet
 from .sstdpn import SSTDPN
 from .syncnet import SyncNet
+from .tcformer import TCFormer
 from .tcn import BDTCN, TCN
 from .tidnet import TIDNet
 from .tsinception import TSception
@@ -130,6 +131,7 @@ __all__ = [
     "SPARCNet",
     "SyncNet",
     "BDTCN",
+    "TCFormer",
     "TCN",
     "TIDNet",
     "TSception",

{braindecode-1.6.1.dev1056 → braindecode-1.6.1.dev1065}/braindecode/models/atcnet.py

@@ -2,14 +2,21 @@
 #
 # License: BSD (3-clause)
 import math
+from typing import Optional
 
 import torch
 from einops.layers.torch import Rearrange
 from mne.utils import warn
 from torch import nn
+from torch.nn.utils.parametrize import is_parametrized, register_parametrization
 
 from braindecode.models.base import EEGModuleMixin
-from braindecode.modules import CausalConv1d, Ensure4d, MaxNormLinear
+from braindecode.modules import (
+    CausalConv1d,
+    Ensure4d,
+    MaxNormLinear,
+    MaxNormParametrize,
+)
 
 
 class ATCNet(EEGModuleMixin, nn.Module):
@@ -213,6 +220,16 @@ class ATCNet(EEGModuleMixin, nn.Module):
     max_norm_const : float
         Maximum L2-norm constraint imposed on weights of the last
         fully-connected layer. Defaults to 0.25.
+    conv_max_norm_const : float | None
+        If not ``None``, applies a max-norm constraint (via a weight
+        parametrization) to the convolution kernels of :class:`_ConvBlock` and
+        :class:`_TCNResidualBlock`, matching the official implementation which
+        uses ``0.6``. When ``None`` (default), no constraint is applied to those
+        layers, preserving the previous behavior. The complementary ``L2``
+        weight decay of the official code can be obtained from
+        :meth:`source_optimizer_param_groups`.
+
+        .. versionadded:: 1.6.1
 
     Notes
     -----
@@ -257,6 +274,7 @@ class ATCNet(EEGModuleMixin, nn.Module):
         tcn_activation: type[nn.Module] = nn.ELU,
         concat=False,
         max_norm_const=0.25,
+        conv_max_norm_const: Optional[float] = None,
         chs_info=None,
         n_times=None,
     ):
@@ -325,6 +343,7 @@ class ATCNet(EEGModuleMixin, nn.Module):
         self.tcn_activation = tcn_activation
         self.concat = concat
         self.max_norm_const = max_norm_const
+        self.conv_max_norm_const = conv_max_norm_const
         self.tcn_n_filters = int(self.conv_block_depth_mult * self.conv_block_n_filters)
         map = dict()
         for w in range(self.n_windows):
@@ -346,6 +365,7 @@ class ATCNet(EEGModuleMixin, nn.Module):
             pool_size_2=conv_block_pool_size_2,
             depth_mult=conv_block_depth_mult,
             dropout=conv_block_dropout,
+            conv_max_norm=self.conv_max_norm_const,
         )
 
         self.F2 = int(conv_block_depth_mult * conv_block_n_filters)
@@ -375,6 +395,7 @@ class ATCNet(EEGModuleMixin, nn.Module):
                             dropout=self.tcn_dropout,
                             activation=self.tcn_activation,
                             dilation=2**i,
+                            conv_max_norm=self.conv_max_norm_const,
                         )
                         for i in range(self.tcn_depth)
                     ]
@@ -458,6 +479,93 @@ class ATCNet(EEGModuleMixin, nn.Module):
 
         return self.out_fun(sw_concat_agg)
 
+    def source_optimizer_param_groups(
+        self,
+        conv_weight_decay: float = 0.009,
+        dense_weight_decay: float = 0.5,
+    ) -> list[dict]:
+        r"""Source-faithful optimizer parameter groups for ATCNet.
+
+        The official Keras implementation applies ``L2`` weight decay only to the
+        convolution / TCN kernels (``conv_weightDecay = 0.009``) and to the final
+        dense layer (``dense_weightDecay = 0.5``); biases, BatchNorm and attention
+        weights are left undecayed. This reproduces the Keras ``L2`` kernel penalty
+        only for optimizers that implement ``weight_decay`` as a coupled L2 term,
+        e.g. :class:`torch.optim.SGD` or :class:`torch.optim.Adam` (as used by the
+        official code). For decoupled-weight-decay optimizers such as
+        :class:`torch.optim.AdamW`, ``weight_decay`` is *not* equivalent to a Keras
+        ``L2`` penalty.
+
+        Parameters
+        ----------
+        conv_weight_decay : float
+            Weight decay applied to the convolution and TCN kernels. Defaults to
+            ``0.009`` as in the official code.
+        dense_weight_decay : float
+            Weight decay applied to the final (dense) layer weights. Defaults to
+            ``0.5`` as in the official code.
+
+        Returns
+        -------
+        list of dict
+            Parameter groups ready to be passed to a ``torch.optim`` optimizer:
+            convolution/TCN kernels, final dense weights, and everything else
+            (no decay).
+
+        Examples
+        --------
+        >>> model = ATCNet(n_chans=22, n_outputs=4, n_times=1125)  # doctest: +SKIP
+        >>> opt = torch.optim.Adam(model.source_optimizer_param_groups())  # doctest: +SKIP
+        """
+
+        def weight_leaf(module):
+            # Return the leaf weight Parameter, accounting for a registered
+            # max-norm parametrization (``conv_max_norm_const`` / MaxNormLinear).
+            if is_parametrized(module, "weight"):
+                return module.parametrizations.weight.original
+            return module.weight
+
+        conv_weights: list[nn.Parameter] = []
+        dense_weights: list[nn.Parameter] = []
+        decayed_ids: set[int] = set()
+
+        # Convolution / TCN kernels (Keras L2 = 0.009).
+        for parent in (self.conv_block, self.temporal_conv_nets):
+            for module in parent.modules():
+                if isinstance(module, (nn.Conv1d, nn.Conv2d)):
+                    weight = weight_leaf(module)
+                    conv_weights.append(weight)
+                    decayed_ids.add(id(weight))
+
+        # Final dense / readout weights (Keras L2 = 0.5).
+        for module in self.final_layer.modules():
+            if isinstance(module, nn.Linear):
+                weight = weight_leaf(module)
+                dense_weights.append(weight)
+                decayed_ids.add(id(weight))
+
+        # Everything else (biases, BatchNorm, attention) is left undecayed.
+        other = [p for p in self.parameters() if id(p) not in decayed_ids]
+
+        return [
+            {"params": conv_weights, "weight_decay": conv_weight_decay},
+            {"params": dense_weights, "weight_decay": dense_weight_decay},
+            {"params": other, "weight_decay": 0.0},
+        ]
+
+
+def _register_max_norm(max_norm, *convs):
+    """Register a max-norm weight parametrization on each given convolution.
+
+    No-op when ``max_norm`` is ``None`` or a passed layer is not a convolution
+    (e.g. an ``nn.Identity`` residual shortcut).
+    """
+    if max_norm is None:
+        return
+    for conv in convs:
+        if isinstance(conv, (nn.Conv1d, nn.Conv2d)):
+            register_parametrization(conv, "weight", MaxNormParametrize(max_norm))
+
 
 class _ConvBlock(nn.Module):
     r"""Convolutional block proposed in ATCNet [1]_, inspired by the EEGNet.
@@ -487,6 +595,7 @@ class _ConvBlock(nn.Module):
         pool_size_2=7,
         depth_mult=2,
         dropout=0.3,
+        conv_max_norm: Optional[float] = None,
     ):
         super().__init__()
 
@@ -534,6 +643,9 @@ class _ConvBlock(nn.Module):
 
         self.drop3 = nn.Dropout2d(dropout)
 
+        # Optional source-faithful max-norm constraint on the conv kernels.
+        _register_max_norm(conv_max_norm, self.conv1, self.conv2, self.conv3)
+
     def forward(self, X):
         # ----- Temporal convolution -----
         # Dimension: (batch_size, 1, T, C)
@@ -662,6 +774,7 @@ class _TCNResidualBlock(nn.Module):
         dropout=0.3,
         activation: type[nn.Module] = nn.ELU,
         dilation=1,
+        conv_max_norm: Optional[float] = None,
     ):
         super().__init__()
         self.activation = activation()
@@ -706,6 +819,10 @@ class _TCNResidualBlock(nn.Module):
         else:
             self.reshaping_conv = nn.Identity()
 
+        # Optional source-faithful max-norm constraint on the conv kernels
+        # (reshaping_conv may be nn.Identity, which the helper skips).
+        _register_max_norm(conv_max_norm, self.conv1, self.conv2, self.reshaping_conv)
+
     def forward(self, X):
         # Dimension: (batch_size, F2, Tw)
         # ----- Double dilated convolutions -----

{braindecode-1.6.1.dev1056 → braindecode-1.6.1.dev1065}/braindecode/models/eegtcnet.py

@@ -9,6 +9,7 @@ import torch.nn as nn
 from einops.layers.torch import Rearrange
 
 from braindecode.models.base import EEGModuleMixin
+from braindecode.models.util import _disable_batch_norm_training_if_batch_size_one
 from braindecode.modules import Chomp1d, MaxNormLinear
 
 
@@ -33,8 +34,10 @@ class EEGTCNet(EEGModuleMixin, nn.Module):
         Number of temporal filters in the first convolutional layer. Default is 8.
     kern_length : int, optional
         Length of the temporal kernel in the first convolutional layer. Default is 64.
-    dropout : float, optional
-        Dropout rate. Default is 0.5.
+    drop_prob : float, optional
+        Default dropout rate, used for both the EEGNet front-end and the TCN
+        block unless overridden by ``drop_prob_eeg`` / ``drop_prob_tcn``.
+        Default is 0.5.
     depth : int, optional
         Number of residual blocks in the TCN. Default is 2.
     kernel_size : int, optional
@@ -44,6 +47,20 @@ class EEGTCNet(EEGModuleMixin, nn.Module):
     max_norm_const : float
         Maximum L2-norm constraint imposed on weights of the last
         fully-connected layer. Defaults to 0.25.
+    drop_prob_eeg : float | None, optional
+        Dropout rate for the EEGNet front-end. If ``None`` (default), falls back
+        to ``drop_prob``. The source/paper uses ``0.2`` for BCI IV 2a.
+    drop_prob_tcn : float | None, optional
+        Dropout rate for the TCN block. If ``None`` (default), falls back to
+        ``drop_prob``. The source/paper uses ``0.3`` for BCI IV 2a.
+    tcn_batch_norm : bool, optional
+        If ``True`` (default), insert a :class:`~torch.nn.BatchNorm1d` after each
+        TCN convolution and use a bias in the residual ``1x1`` downsample, matching
+        the original Keras source. Set to ``False`` to recover the previous
+        braindecode architecture (e.g. to load checkpoints trained before this
+        change). Default is ``True``.
+
+        .. versionadded:: 1.6.1
 
     References
     ----------
@@ -72,6 +89,9 @@ class EEGTCNet(EEGModuleMixin, nn.Module):
         kernel_size: int = 4,
         filters: int = 12,
         max_norm_const: float = 0.25,
+        drop_prob_eeg: float | None = None,
+        drop_prob_tcn: float | None = None,
+        tcn_batch_norm: bool = True,
     ):
         super().__init__(
             n_outputs=n_outputs,
@@ -85,6 +105,11 @@ class EEGTCNet(EEGModuleMixin, nn.Module):
 
         self.activation = activation
         self.drop_prob = drop_prob
+        # Separate dropout rates for the EEGNet front-end and the TCN block;
+        # fall back to ``drop_prob`` when not explicitly set.
+        self.drop_prob_eeg = drop_prob if drop_prob_eeg is None else drop_prob_eeg
+        self.drop_prob_tcn = drop_prob if drop_prob_tcn is None else drop_prob_tcn
+        self.tcn_batch_norm = tcn_batch_norm
         self.depth_multiplier = depth_multiplier
         self.filter_1 = filter_1
         self.kern_length = kern_length
@@ -103,7 +128,7 @@ class EEGTCNet(EEGModuleMixin, nn.Module):
             filter_1=self.filter_1,
             kern_length=self.kern_length,
             depth_multiplier=self.depth_multiplier,
-            drop_prob=self.drop_prob,
+            drop_prob=self.drop_prob_eeg,
             activation=self.activation,
         )
         self.arrange_dim_eegnet = Rearrange(
@@ -116,8 +141,9 @@ class EEGTCNet(EEGModuleMixin, nn.Module):
             depth=self.depth,
             kernel_size=self.kernel_size,
             filters=self.filters,
-            drop_prob=self.drop_prob,
+            drop_prob=self.drop_prob_tcn,
             activation=self.activation,
+            batch_norm=self.tcn_batch_norm,
         )
 
         # Classification Block
@@ -127,6 +153,7 @@ class EEGTCNet(EEGModuleMixin, nn.Module):
             max_norm_val=self.max_norm_const,
         )
 
+    @_disable_batch_norm_training_if_batch_size_one
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Forward pass of the EEGTCNet model.
@@ -277,6 +304,7 @@ class _TCNBlock(nn.Module):
         filters: int,
         drop_prob: float,
         activation: type[nn.Module] = nn.ELU,
+        batch_norm: bool = True,
     ):
         super().__init__()
         self.activation = activation()
@@ -284,10 +312,15 @@ class _TCNBlock(nn.Module):
         self.depth = depth
         self.filters = filters
         self.kernel_size = kernel_size
+        self.batch_norm = batch_norm
 
         self.layers = nn.ModuleList()
+        # The Keras source uses a bias on the residual 1x1 downsample, paired
+        # with the per-conv BatchNorm; the previous braindecode default omitted
+        # both. Keeping them on the same flag makes ``batch_norm=False`` reproduce
+        # the legacy architecture exactly.
         self.downsample = (
-            nn.Conv1d(input_dimension, filters, kernel_size=1, bias=False)
+            nn.Conv1d(input_dimension, filters, kernel_size=1, bias=batch_norm)
             if input_dimension != filters
             else None
         )
@@ -295,32 +328,34 @@ class _TCNBlock(nn.Module):
         for i in range(depth):
             dilation = 2**i
             padding = (kernel_size - 1) * dilation
+            in_dim = input_dimension if i == 0 else filters
             conv_block = nn.Sequential(
-                nn.Conv1d(
-                    in_channels=input_dimension if i == 0 else filters,
-                    out_channels=filters,
-                    kernel_size=kernel_size,
-                    dilation=dilation,
-                    padding=padding,
-                    bias=False,
-                ),
-                Chomp1d(padding),
-                self.activation,
-                nn.Dropout(self.drop_prob),
-                nn.Conv1d(
-                    in_channels=filters,
-                    out_channels=filters,
-                    kernel_size=kernel_size,
-                    dilation=dilation,
-                    padding=padding,
-                    bias=False,
-                ),
-                Chomp1d(padding),
-                self.activation,
-                nn.Dropout(self.drop_prob),
+                *self._conv_unit(in_dim, dilation, padding),
+                *self._conv_unit(filters, dilation, padding),
             )
             self.layers.append(conv_block)
 
+    def _conv_unit(self, in_channels: int, dilation: int, padding: int):
+        # Conv -> Chomp -> [BatchNorm] -> activation -> Dropout, as in the
+        # original EEG-TCNet Keras implementation.
+        unit = [
+            nn.Conv1d(
+                in_channels=in_channels,
+                out_channels=self.filters,
+                kernel_size=self.kernel_size,
+                dilation=dilation,
+                padding=padding,
+                bias=False,
+            ),
+            Chomp1d(padding),
+        ]
+        if self.batch_norm:
+            # Keras BatchNormalization defaults: momentum=0.99 (-> 0.01 in
+            # PyTorch convention) and epsilon=1e-3.
+            unit.append(nn.BatchNorm1d(self.filters, momentum=0.01, eps=1e-3))
+        unit += [self.activation, nn.Dropout(self.drop_prob)]
+        return unit
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # x shape: (batch_size, time_steps, input_dimension)
         x = x.permute(0, 2, 1)  # (batch_size, input_dimension, time_steps)

{braindecode-1.6.1.dev1056 → braindecode-1.6.1.dev1065}/braindecode/models/sstdpn.py

@@ -183,7 +183,11 @@ class SSTDPN(EEGModuleMixin, nn.Module):
         This constraint acts as an implicit force to push features away from the origin. Default is 1.0.
 
     proto_cpt_std : float, optional
-        Standard deviation for Intra-class Compactness Prototype initialization. Default is 0.01.
+        Standard deviation for Intra-class Compactness Prototype initialization. Default is 1.0,
+        matching the source ``torch.randn`` initialization.
+
+        .. versionchanged:: 1.6.1
+            Default changed from ``0.01`` to ``1.0`` to match the original implementation.
 
     spt_attn_global_context_kernel : int, optional
         Kernel size for global context embedding in Spatial-Spectral Attention module.
@@ -231,7 +235,7 @@ class SSTDPN(EEGModuleMixin, nn.Module):
         mvp_kernel_sizes: Optional[List[int]] = None,
         return_features: bool = False,
         proto_sep_maxnorm: float = 1.0,
-        proto_cpt_std: float = 0.01,
+        proto_cpt_std: float = 1.0,
         spt_attn_global_context_kernel: int = 250,
         spt_attn_epsilon: float = 1e-5,
         spt_attn_mode: str = "var",
@@ -327,9 +331,11 @@ class SSTDPN(EEGModuleMixin, nn.Module):
         """
 
         features = self.encoder(x)  # (b, feat_dim)
-        # Renormalize inter-class separation prototypes
+        # Renormalize inter-class separation prototypes. ``proto_sep`` has shape
+        # (n_outputs, feat_dim); the source constrains each class-row prototype
+        # vector (``||s_i|| <= S``), i.e. renorm along dim=0.
         self.proto_sep.data = torch.renorm(
-            self.proto_sep.data, p=2, dim=1, maxnorm=self.proto_sep_maxnorm
+            self.proto_sep.data, p=2, dim=0, maxnorm=self.proto_sep_maxnorm
         )
         logits = torch.einsum("bd,cd->bc", features, self.proto_sep)  # (b, n_outputs)
         logits = self.final_layer(logits)

{braindecode-1.6.1.dev1056 → braindecode-1.6.1.dev1065}/braindecode/models/summary.csv

@@ -56,4 +56,5 @@ InterpolatedBIOT,"Sleep Staging, Epilepsy",Classification,200,"chs_info, n_outpu
 InterpolatedEEGPT,General,Classification,256,"chs_info, n_outputs, n_times",25323137,"InterpolatedEEGPT(chs_info=<user>, n_outputs=4, n_times=1024)","Foundation Model,Attention/Transformer,Channel"
 InterpolatedLaBraM,General,"Classification, Embedding",200,"chs_info, n_outputs, n_times",5866196,"InterpolatedLaBraM(chs_info=<user>, n_outputs=4, n_times=1000)","Convolution,Foundation Model,Channel"
 InterpolatedSignalJEPA,"Motor Imagery, ERP, SSVEP",Embedding,128,"chs_info",3456882,"InterpolatedSignalJEPA(chs_info=<user>)","Convolution,Channel,Foundation Model"
+TCFormer,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times",77820,"TCFormer(n_chans=22, n_outputs=4, n_times=1000)","Convolution,Attention/Transformer"
 EEGDINO,General,Classification,200,"n_chans, n_outputs, n_times",4539698,"EEGDINO(n_chans=19, n_outputs=4, n_times=1000)","Foundation Model,Attention/Transformer"