braindecode 1.6.1.dev1065__tar.gz → 1.6.1.dev1073__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.6.1.dev1065
+Version: 1.6.1.dev1073
 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.dev1065 → braindecode-1.6.1.dev1073}/braindecode/augmentation/base.py

@@ -105,7 +105,15 @@ class Transform(torch.nn.Module):
         mask = self._get_mask(out_X.shape[0], out_X.device)
         num_valid = mask.sum().long()
 
-        if num_valid > 0:
+        if num_valid == out_X.shape[0]:
+            # Fast path: every sample is augmented (e.g. ``probability=1.0``,
+            # the common case). Operate on the whole batch and skip the boolean
+            # gather/scatter (``out_X[mask]`` then ``out_X[mask] = ...``), which
+            # would otherwise copy the batch twice.
+            out_X, out_y = self.operation(
+                out_X, out_y, **self.get_augmentation_params(out_X, out_y)
+            )
+        elif num_valid > 0:
             # Uses the mask to define the output
             out_X[mask, ...], tr_y = self.operation(
                 out_X[mask, ...],
@@ -168,23 +176,47 @@ class Compose(Transform):
         return X, y
 
 
-def _make_collateable(transform, device=None):
-    """Wraps a transform to make it collateable.
-
-    with device control.
-    """
+class _AugmentationCollate:
+    """Collate that applies a transform to each batch, with optional expansion.
 
-    def _collate_fn(batch):
-        collated_batch = default_collate(batch)
-        X, y = collated_batch[:2]
+    A picklable callable (so it works with ``num_workers > 0``, unlike a
+    closure) that collates the batch once and then either augments it in place
+    or grows it.
 
-        if device is not None:
-            X = X.to(device)
-            y = y.to(device)
-
-        return (*transform(X, y), *collated_batch[2:])
+    Parameters
+    ----------
+    transform : Transform
+        Transform applied to the collated ``(X, y)``.
+    device : str | torch.device | None, optional
+        Device the batch is moved to before transforming. Defaults to None.
+    n_augmentation : int, optional
+        ``0`` (default) applies the transform in place (batch size unchanged).
+        ``> 0`` keeps the clean originals and appends ``n_augmentation``
+        independently transformed copies, returning ``(X, y)`` of
+        ``(1 + n_augmentation)`` times the original size.
+    """
 
-    return _collate_fn
+    def __init__(self, transform, device=None, n_augmentation=0):
+        self.transform = transform
+        self.device = device
+        self.n_augmentation = n_augmentation
+
+    def __call__(self, batch):
+        collated = default_collate(batch)
+        X, y = collated[:2]
+        if self.device is not None:
+            X = X.to(self.device)
+            y = y.to(self.device)
+        if self.n_augmentation == 0:
+            return (*self.transform(X, y), *collated[2:])
+        # Fixed expansion: keep the clean originals, then append
+        # ``n_augmentation`` independently transformed copies.
+        xs, ys = [X], [y]
+        for _ in range(self.n_augmentation):
+            aug_X, aug_y = self.transform(X, y)
+            xs.append(aug_X)
+            ys.append(aug_y)
+        return torch.cat(xs), torch.cat(ys)
 
 
 class AugmentedDataLoader(DataLoader):
@@ -198,28 +230,44 @@ class AugmentedDataLoader(DataLoader):
         Transform or sequence of Transform to be applied to each batch.
     device : str | torch.device | None, optional
         Device on which to transform the data. Defaults to None.
+    n_augmentation : int, optional
+        Number of augmented copies to append to each batch (fixed expansion).
+        When ``0`` (default) the transforms are applied in place and the batch
+        keeps its size (stochastic-per-epoch augmentation, backwards-compatible).
+        When ``> 0`` each batch becomes ``(1 + n_augmentation)`` times larger:
+        the clean originals are kept and ``n_augmentation`` independently
+        transformed copies are appended. This expresses augmentations defined as
+        a fixed set-expansion (e.g. the EEG-Inception MI 6x training set,
+        ``n_augmentation=5``). In this mode batches are returned as ``(X, y)``.
     **kwargs : dict, optional
         keyword arguments to pass to standard DataLoader class.
     """
 
-    def __init__(self, dataset, transforms=None, device=None, **kwargs):
+    def __init__(
+        self, dataset, transforms=None, device=None, n_augmentation=0, **kwargs
+    ):
         if "collate_fn" in kwargs:
             raise ValueError(
                 "collate_fn cannot be used in this context because it is used "
                 "to pass transform"
             )
+        if n_augmentation < 0:
+            raise ValueError("n_augmentation must be a non-negative integer.")
         if transforms is None or (
             isinstance(transforms, list) and len(transforms) == 0
         ):
-            self.collated_tr = _make_collateable(IdentityTransform(), device=device)
+            transform = IdentityTransform()
         elif isinstance(transforms, (Transform, nn.Module)):
-            self.collated_tr = _make_collateable(transforms, device=device)
+            transform = transforms
         elif isinstance(transforms, list):
-            self.collated_tr = _make_collateable(Compose(transforms), device=device)
+            transform = Compose(transforms)
         else:
             raise TypeError(
                 "transforms can be either a Transform object "
                 "or a list of Transform objects."
             )
 
+        self.collated_tr = _AugmentationCollate(
+            transform, device=device, n_augmentation=n_augmentation
+        )
         super().__init__(dataset, collate_fn=self.collated_tr, **kwargs)

{braindecode-1.6.1.dev1065 → braindecode-1.6.1.dev1073}/braindecode/augmentation/functional.py

@@ -18,7 +18,7 @@ from mne.filter import notch_filter
 from scipy.interpolate import Rbf
 from sklearn.utils import check_random_state
 from torch.fft import fft, ifft
-from torch.nn.functional import one_hot, pad
+from torch.nn.functional import pad
 
 
 def identity(X: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
@@ -253,25 +253,34 @@ def channels_dropout(
     return X * mask.unsqueeze(-1), y
 
 
-def _make_permutation_matrix(
+def _make_channel_permutation(
     X: torch.Tensor, mask: torch.Tensor, random_state: int | np.random.Generator | None
 ) -> torch.Tensor:
+    """Per-sample channel permutation that shuffles only the masked channels.
+
+    Returns a ``(batch, n_channels)`` index permutation: output channel ``i`` is
+    taken from input channel ``perm[b, i]``. Masked channels are permuted among
+    their own positions; unmasked channels keep their place.
+
+    The per-sample ``rng.permutation`` is kept (rather than a vectorized random
+    draw) so the output is bit-for-bit identical to the previous one-hot/matmul
+    implementation — vectorizing the draw would change seeded results. The win
+    is downstream: the caller gathers with this index instead of building
+    ``(batch, n_channels, n_channels)`` one-hot matrices and a matmul.
+    """
     rng = check_random_state(random_state)
     batch_size, n_channels, _ = X.shape
-    hard_mask = mask.round()
-    batch_permutations = torch.empty(
-        batch_size, n_channels, n_channels, device=X.device
-    )
-    for b, mask in enumerate(hard_mask):
-        channels_to_shuffle = torch.arange(n_channels, device=X.device)
-        channels_to_shuffle = channels_to_shuffle[mask.bool()]
-        reordered_channels = torch.tensor(
-            rng.permutation(channels_to_shuffle.cpu()), device=X.device
+    hard_mask = mask.round().bool()
+    channels = torch.arange(n_channels, device=X.device)
+    perm = channels.repeat(batch_size, 1)
+    for b in range(batch_size):
+        channels_to_shuffle = channels[hard_mask[b]]
+        perm[b, channels_to_shuffle] = torch.tensor(
+            rng.permutation(channels_to_shuffle.cpu()),
+            device=X.device,
+            dtype=perm.dtype,
         )
-        channels_permutation = torch.arange(n_channels, device=X.device)
-        channels_permutation[channels_to_shuffle] = reordered_channels
-        batch_permutations[b, ...] = one_hot(channels_permutation)
-    return batch_permutations
+    return perm
 
 
 def channels_shuffle(
@@ -314,8 +323,11 @@ def channels_shuffle(
     if p_shuffle == 0:
         return X, y
     mask = _pick_channels_randomly(X, 1 - p_shuffle, random_state)
-    batch_permutations = _make_permutation_matrix(X, mask, random_state)
-    return torch.matmul(batch_permutations, X), y
+    perm = _make_channel_permutation(X, mask, random_state)
+    # Gather channels (output[b, i] = X[b, perm[b, i]]) instead of building
+    # one-hot permutation matrices and a (batch, C, C) @ (batch, C, T) matmul.
+    perm = perm.unsqueeze(-1).expand(-1, -1, X.shape[-1])
+    return torch.gather(X, 1, perm), y
 
 
 def gaussian_noise(
@@ -1188,19 +1200,21 @@ def mask_encoding(
        32 (2024): 875-886.
     """
 
-    batch_indices = torch.arange(X.shape[0]).repeat_interleave(n_segments)
-    start_indices = time_start.flatten()
-    mask_indices = start_indices[:, None] + torch.arange(segment_length)
-
-    # Create a boolean mask with the same shape as X
-    mask = torch.zeros_like(X, dtype=torch.bool)
-    for batch_index, grouped_mask_indices in zip(batch_indices, mask_indices):
-        mask[batch_index, :, grouped_mask_indices] = True
-
-    # Apply the mask to set the values to 0
-    X[mask] = 0
-
-    return X, y  # Return the masked tensor and labels
+    # All channels are zeroed at the same time positions, so build a single
+    # (batch, n_times) time-mask and broadcast it over channels, vectorized
+    # (no Python loop over batch * n_segments).
+    device = X.device
+    batch_indices = torch.arange(X.shape[0], device=device).repeat_interleave(
+        n_segments * segment_length
+    )
+    seg_indices = (
+        time_start.flatten().to(device)[:, None]
+        + torch.arange(segment_length, device=device)
+    ).flatten()
+    time_mask = torch.zeros(X.shape[0], X.shape[-1], dtype=torch.bool, device=device)
+    time_mask[batch_indices, seg_indices] = True
+
+    return X.masked_fill(time_mask.unsqueeze(1), 0), y
 
 
 def channels_rereference(

{braindecode-1.6.1.dev1065 → braindecode-1.6.1.dev1073}/braindecode/datasets/bids/hub.py

@@ -327,6 +327,7 @@ class HubDatasetMixin:
 
         n_recordings = len(self.datasets)
         first_ds = self.datasets[0]
+        descriptions = [ds.description for ds in self.datasets]
 
         # Get dataset-specific info based on type using registry
         dataset_type = get_dataset_type(first_ds)
@@ -358,6 +359,9 @@ class HubDatasetMixin:
             raise TypeError(f"Unsupported dataset type: {dataset_type}")
 
         # Create README content and save
+        license_value, license_needs_update = _infer_license_from_descriptions(
+            descriptions
+        )
         readme_content = _generate_readme_content(
             format_info=format_info,
             n_recordings=n_recordings,
@@ -366,6 +370,8 @@ class HubDatasetMixin:
             data_type=data_type,
             n_windows=n_windows,
             total_duration=total_duration,
+            license=license_value,
+            license_needs_update=license_needs_update,
         )
 
         # Save README
@@ -967,6 +973,8 @@ def _generate_readme_content(
     n_windows: int,
     total_duration: float | None = None,
     format: str = "zarr",
+    license: str = "please-specify",
+    license_needs_update: bool = True,
 ):
     """Generate README.md content for a dataset uploaded to the Hub.
 
@@ -988,6 +996,10 @@ def _generate_readme_content(
         Total duration in seconds across all recordings.
     format : str
         Storage format (default: "zarr").
+    license : str
+        Dataset license string to write in YAML front matter.
+    license_needs_update : bool
+        Whether to append a reminder comment next to the license field.
 
     Returns
     -------
@@ -1003,6 +1015,10 @@ def _generate_readme_content(
         str(timedelta(seconds=int(total_duration))) if total_duration else "N/A"
     )
 
+    license_line = f"license: {license}"
+    if license_needs_update:
+        license_line += "  # Please update this field to reflect your dataset's license"
+
     return f"""---
 tags:
 - braindecode
@@ -1010,7 +1026,7 @@ tags:
 - neuroscience
 - brain-computer-interface
 - deep-learning
-license: unknown
+{license_line}
 ---
 
 # EEG Dataset
@@ -1106,3 +1122,27 @@ if hasattr(dataset.datasets[0], "bids_channels"):
 
 *Created with [braindecode](https://braindecode.org)*
 """
+
+
+def _infer_license_from_descriptions(
+    descriptions: list[pd.Series | None],
+) -> tuple[str, bool]:
+    """Infer a single dataset license from descriptions.
+
+    Returns inferred license and whether a manual update reminder is needed.
+    """
+    licenses = set()
+    for description in descriptions:
+        if description is None:
+            continue
+        raw_license = description.get("license", None)
+        if raw_license is None:
+            continue
+        normalized = str(raw_license).strip()
+        if normalized and normalized.lower() != "nan":
+            licenses.add(normalized)
+
+    if len(licenses) == 1:
+        return next(iter(licenses)), False
+
+    return "please-specify", True

{braindecode-1.6.1.dev1065 → braindecode-1.6.1.dev1073}/braindecode/models/__init__.py

@@ -54,6 +54,7 @@ from .sleep_stager_blanco_2020 import SleepStagerBlanco2020
 from .sleep_stager_chambon_2018 import SleepStagerChambon2018
 from .sparcnet import SPARCNet
 from .sstdpn import SSTDPN
+from .steegformer import STEEGFormer
 from .syncnet import SyncNet
 from .tcformer import TCFormer
 from .tcn import BDTCN, TCN
@@ -129,6 +130,7 @@ __all__ = [
     "SleepStagerBlanco2020",
     "SleepStagerChambon2018",
     "SPARCNet",
+    "STEEGFormer",
     "SyncNet",
     "BDTCN",
     "TCFormer",

{braindecode-1.6.1.dev1065 → braindecode-1.6.1.dev1073}/braindecode/models/eegitnet.py

@@ -4,9 +4,16 @@
 # License: BSD-3
 from einops.layers.torch import Rearrange
 from torch import nn
+from torch.nn.utils.parametrize import register_parametrization
 
 from braindecode.models.base import EEGModuleMixin
-from braindecode.modules import DepthwiseConv2d, Ensure4d, InceptionBlock
+from braindecode.modules import (
+    DepthwiseConv2d,
+    Ensure4d,
+    InceptionBlock,
+    LinearWithConstraint,
+    MaxNormParametrize,
+)
 
 
 class EEGITNet(EEGModuleMixin, nn.Sequential):
@@ -95,7 +102,9 @@ class EEGITNet(EEGModuleMixin, nn.Sequential):
             sfreq=sfreq,
         )
         self.mapping = {
-            "classification.1.weight": "final_layer.weight",
+            # final_layer is max-norm constrained, so its weight lives under the
+            # parametrization buffer rather than ``final_layer.weight``.
+            "classification.1.weight": "final_layer.parametrizations.weight.original",
             "classification.1.bias": "final_layer.bias",
         }
 
@@ -185,8 +194,10 @@ class EEGITNet(EEGModuleMixin, nn.Sequential):
         self.add_module(
             "dim_reduction",
             nn.Sequential(
-                nn.Conv2d(tcn_in_channel, tcn_in_channel * 2, kernel_size=(1, 1)),
-                nn.BatchNorm2d(tcn_in_channel * 2),
+                # The authors' DR layer keeps the width (14 filters, paper III-C),
+                # it does not double it.
+                nn.Conv2d(tcn_in_channel, tcn_in_channel, kernel_size=(1, 1)),
+                nn.BatchNorm2d(tcn_in_channel),
                 activation(),
                 nn.AvgPool2d((1, tcn_kernel_size)),
                 nn.Dropout(drop_prob),
@@ -198,7 +209,11 @@ class EEGITNet(EEGModuleMixin, nn.Sequential):
 
         num_features = self.get_output_shape()[-1]
 
-        self.add_module("final_layer", nn.Linear(num_features, self.n_outputs))
+        # The authors constrain the final dense with ``max_norm(0.25)``.
+        self.add_module(
+            "final_layer",
+            LinearWithConstraint(num_features, self.n_outputs, max_norm=0.25),
+        )
 
     @staticmethod
     def _get_inception_branch(
@@ -208,6 +223,18 @@ class EEGITNet(EEGModuleMixin, nn.Sequential):
         depth_multiplier=1,
         activation: type[nn.Module] = nn.ELU,
     ):
+        # The authors constrain the spatial depthwise conv with
+        # ``depthwise_constraint = MaxNorm(max_value=1)``.
+        spatial_depthwise = DepthwiseConv2d(
+            out_channels,
+            kernel_size=(in_channels, 1),
+            depth_multiplier=depth_multiplier,
+            bias=False,
+            padding="valid",
+        )
+        register_parametrization(
+            spatial_depthwise, "weight", MaxNormParametrize(max_norm=1.0)
+        )
         return nn.Sequential(
             nn.Conv2d(
                 1,
@@ -217,13 +244,7 @@ class EEGITNet(EEGModuleMixin, nn.Sequential):
                 bias=False,
             ),
             nn.BatchNorm2d(out_channels),
-            DepthwiseConv2d(
-                out_channels,
-                kernel_size=(in_channels, 1),
-                depth_multiplier=depth_multiplier,
-                bias=False,
-                padding="valid",
-            ),
+            spatial_depthwise,
             nn.BatchNorm2d(out_channels),
             activation(),
         )

{braindecode-1.6.1.dev1065 → braindecode-1.6.1.dev1073}/braindecode/models/eegsym.py

@@ -504,6 +504,11 @@ class _InceptionBlock(nn.Module):
                             out_channels=filters_per_branch * len(scales_samples),
                             kernel_size=(1, 1, ncha),
                             padding=(0, 0, 0),
+                            # depthwise/grouped spatial conv with no bias, as in
+                            # the authors' ``unit_dconv`` (per-filter spatial
+                            # filtering, not dense mixing across filters).
+                            groups=filters_per_branch * len(scales_samples),
+                            bias=False,
                         ),
                         nn.BatchNorm3d(filters_per_branch * len(scales_samples)),
                         activation,
@@ -621,6 +626,10 @@ class _ResidualBlock(nn.Module):
                             out_channels=filters,
                             kernel_size=(1, 1, ncha),  # Spatial convolution
                             padding=(0, 0, 0),
+                            # depthwise/grouped spatial conv with no bias, as in
+                            # the authors' ``unit_dconv``.
+                            groups=filters,
+                            bias=False,
                         ),
                         nn.BatchNorm3d(filters),
                         activation,

{braindecode-1.6.1.dev1065 → braindecode-1.6.1.dev1073}/braindecode/models/msvtnet.py

@@ -186,9 +186,13 @@ class MSVTNet(EEGModuleMixin, nn.Module):
 
         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 BOTH the main-head logits and the stacked branch logits, as
+            # in the authors' implementation (``return x, bx``). This is required
+            # for the paper's joint deep-supervision loss
+            # ``lambda * CE(main) + (1 - lambda) * sum_i CE(branch_i)`` and lets
+            # callers score on the main head.
+            # x: [batch_size, n_classes]; branches: [n_branches, batch_size, n_classes]
+            return x, torch.stack(branch_preds)
         return x