braindecode 1.2.0.dev184328194__py3-none-any.whl → 1.3.0.dev171178473__py3-none-any.whl

braindecode/models/attentionbasenet.py

@@ -26,24 +26,151 @@ from braindecode.modules.attention import (
 class AttentionBaseNet(EEGModuleMixin, nn.Module):
     """AttentionBaseNet from Wimpff M et al. (2023) [Martin2023]_.
 
+    :bdg-success:`Convolution` :bdg-info:`Small Attention`
+
     .. figure:: https://content.cld.iop.org/journals/1741-2552/21/3/036020/revision2/jnead48b9f2_hr.jpg
-       :align: center
-       :alt: Attention Base Net
+        :align: center
+        :alt: AttentionBaseNet Architecture
+        :width: 640px
+
+
+    .. rubric:: Architectural Overview
+
+    AttentionBaseNet is a *convolution-first* network with a *channel-attention* stage.
+    The end-to-end flow is:
+
+    - (i) :class:`_FeatureExtractor` learns a temporal filter bank and per-filter spatial
+      projections (depthwise across electrodes), then condenses time by pooling;
+    - (ii) **Channel Expansion** uses a ``1x1`` convolution to set the feature width;
+    - (iii) :class:`_ChannelAttentionBlock` refines features via depthwise–pointwise temporal
+      convs and an optional channel-attention module (SE/CBAM/ECA/…);
+    - (iv) **Classifier** flattens the sequence and applies a linear readout.
+
+    This design mirrors shallow CNN pipelines (EEGNet-style stem) but inserts a pluggable
+    attention unit that *re-weights channels* (and optionally temporal positions) before
+    classification.
+
+
+    .. rubric:: Macro Components
+
+    - :class:`_FeatureExtractor` **(Shallow conv stem → condensed feature map)**
+
+        - *Operations.*
+        - **Temporal conv** (:class:`torch.nn.Conv2d`) with kernel ``(1, L_t)`` creates a learned
+          FIR-like filter bank with ``n_temporal_filters`` maps.
+        - **Depthwise spatial conv** (:class:`torch.nn.Conv2d`, ``groups=n_temporal_filters``)
+          with kernel ``(n_chans, 1)`` learns per-filter spatial projections over the full montage.
+        - **BatchNorm → ELU → AvgPool → Dropout** stabilize and downsample time.
+        - Output shape: ``(B, F2, 1, T₁)`` with ``F2 = n_temporal_filters x spatial_expansion``.
+
+    *Interpretability/robustness.* Temporal kernels behave as analyzable FIR filters; the
+    depthwise spatial step yields rhythm-specific topographies. Pooling acts as a local
+    integrator that reduces variance on short EEG windows.
+
+    - **Channel Expansion**
+
+        - *Operations.*
+        - A ``1x1`` conv → BN → activation maps ``F2 → ch_dim`` without changing
+          the temporal length ``T₁`` (shape: ``(B, ch_dim, 1, T₁)``).
+          This sets the embedding width for the attention block.
+
+    - :class:`_ChannelAttentionBlock` **(temporal refinement + channel attention)**
+
+        - *Operations.*
+        - **Depthwise temporal conv** ``(1, L_a)`` (groups=``ch_dim``) + **pointwise ``1x1``**,
+          BN and activation → preserves shape ``(B, ch_dim, 1, T₁)`` while refining timing.
+        - **Optional attention module** (see *Additional Mechanisms*) applies channel reweighting
+          (some variants also apply temporal gating).
+        - **AvgPool (1, P₂)** with stride ``(1, S₂)`` and **Dropout** → outputs
+          ``(B, ch_dim, 1, T₂)``.
+
+    *Role.* Emphasizes informative channels (and, in certain modes, salient time steps)
+    before the classifier; complements the convolutional priors with adaptive re-weighting.
+
+    - **Classifier (aggregation + readout)**
+
+    *Operations.* :class:`torch.nn.Flatten` → :class:`torch.nn.Linear` from
+    ``(B, ch_dim·T₂)`` to classes.
+
+
+    .. rubric:: Convolutional Details
 
-    Neural Network from the paper: EEG motor imagery decoding:
-    A framework for comparative analysis with channel attention
-    mechanisms
+    - **Temporal (where time-domain patterns are learned).**
+        Wide kernels in the stem (``(1, L_t)``) act as a learned filter bank for oscillatory
+        bands/transients; the attention block's depthwise temporal conv (``(1, L_a)``) sharpens
+        short-term dynamics after downsampling. Pool sizes/strides (``P₁,S₁`` then ``P₂,S₂``)
+        set the token rate and effective temporal resolution.
 
-    The paper and original code with more details about the methodological
-    choices are available at the [Martin2023]_ and [MartinCode]_.
+    - **Spatial (how electrodes are processed).**
+        A depthwise spatial conv with kernel ``(n_chans, 1)`` spans the full montage to
+        learn *per-temporal-filter* spatial projections (no cross-filter mixing at this step),
+        mirroring the interpretable spatial stage in shallow CNNs.
 
-    The AttentionBaseNet architecture is composed of four modules:
-    - Input Block that performs a temporal convolution and a spatial
-    convolution.
-    - Channel Expansion that modifies the number of channels.
-    - An attention block that performs channel attention with several
-    options
-    - ClassificationHead
+    - **Spectral (how frequency content is captured).**
+        No explicit Fourier/wavelet transform is used in the stem—spectral selectivity
+        emerges from learned temporal kernels. When ``attention_mode="fca"``, a frequency
+        channel attention (DCT-based) summarizes frequencies to drive channel weights.
+
+
+    .. rubric:: Attention / Sequential Modules
+
+    - **Type.** Channel attention chosen by ``attention_mode`` (SE, ECA, CBAM, CAT, GSoP,
+        EncNet, GE, GCT, SRM, CATLite). Most operate purely on channels; CBAM/CAT additionally
+        include temporal attention.
+
+    - **Shapes.** Input/Output around attention: ``(B, ch_dim, 1, T₁)``. Re-arrangements
+        (if any) are internal to the module; the block returns the same shape before pooling.
+
+    - **Role.** Re-weights channels (and optionally time) to highlight informative sources
+        and suppress distractors, improving SNR ahead of the linear head.
+
+
+    .. rubric:: Additional Mechanisms
+
+    **Attention variants at a glance:**
+
+    - ``"se"``: Squeeze-and-Excitation (global pooling → bottleneck → gates).
+    - ``"gsop"``: Global second-order pooling (covariance-aware channel weights).
+    - ``"fca"``: Frequency Channel Attention (DCT summary; uses ``seq_len`` and ``freq_idx``).
+    - ``"encnet"``: EncNet with learned codewords (uses ``n_codewords``).
+    - ``"eca"``: Efficient Channel Attention (local 1-D conv over channel descriptor; uses ``kernel_size``).
+    - ``"ge"``: Gather–Excite (context pooling with optional MLP; can use ``extra_params``).
+    - ``"gct"``: Gated Channel Transformation (global context normalization + gating).
+    - ``"srm"``: Style-based recalibration (mean–std descriptors; optional MLP).
+    - ``"cbam"``: Channel then temporal attention (uses ``kernel_size``).
+    - ``"cat"`` / ``"catlite"``: Collaborative (channel ± temporal) attention; *lite* omits temporal.
+
+    **Auto-compatibility on short inputs:**
+
+        If the input duration is too short for the configured kernels/pools, the implementation
+        **automatically rescales** temporal lengths/strides downward (with a warning) to keep
+        shapes valid and preserve the pipeline semantics.
+
+    .. rubric:: Usage and Configuration
+
+    - ``n_temporal_filters``, ``temporal_filter_length`` and ``spatial_expansion``:
+        control the capacity and the number of spatial projections in the stem.
+    - ``pool_length_inp``, ``pool_stride_inp`` then ``pool_length``, ``pool_stride``:
+        trade temporal resolution for compute; they determine the final sequence length ``T₂``.
+    - ``ch_dim``: width after the ``1x1`` expansion and the effective embedding size for attention.
+    - ``attention_mode`` + its specific hyperparameters (``reduction_rate``,
+        ``kernel_size``, ``seq_len``, ``freq_idx``, ``n_codewords``, ``use_mlp``):
+        select and tune the reweighting mechanism.
+    - ``drop_prob_inp`` and ``drop_prob_attn``: regularize stem and attention stages.
+    - **Training tips.**
+
+        Start with moderate pooling (e.g., ``P₁=75,S₁=15``) and ELU activations; enable attention
+        only after the stem learns stable filters. For small datasets, prefer simpler modes
+        (``"se"``, ``"eca"``) before heavier ones (``"gsop"``, ``"encnet"``).
+
+    Notes
+    -----
+    - Sequence length after each stage is computed internally; the final classifier expects
+      a flattened ``ch_dim x T₂`` vector.
+    - Attention operates on *channel* dimension by design; temporal gating exists only in
+      specific variants (CBAM/CAT).
+    - The paper and original code with more details about the methodological
+      choices are available at the [Martin2023]_ and [MartinCode]_.
 
     .. versionadded:: 0.9
 
@@ -73,6 +200,7 @@ class AttentionBaseNet(EEGModuleMixin, nn.Module):
         the depth of the network after the initial layer. Default is 16.
     attention_mode : str, optional
         The type of attention mechanism to apply. If `None`, no attention is applied.
+
         - "se" for Squeeze-and-excitation network
         - "gsop" for Global Second-Order Pooling
         - "fca" for Frequency Channel Attention Network
@@ -83,9 +211,10 @@ class AttentionBaseNet(EEGModuleMixin, nn.Module):
         - "srm" for Style-based Recalibration Module
         - "cbam" for Convolutional Block Attention Module
         - "cat" for Learning to collaborate channel and temporal attention
-        from multi-information fusion
+          from multi-information fusion
         - "catlite" for Learning to collaborate channel attention
-        from multi-information fusion (lite version, cat w/o temporal attention)
+          from multi-information fusion (lite version, cat w/o temporal attention)
+
     pool_length : int, default=8
         The length of the window for the average pooling operation.
     pool_stride : int, default=8
@@ -256,6 +385,8 @@ class AttentionBaseNet(EEGModuleMixin, nn.Module):
         for k, pl, ps in zip(kernel_lengths, pool_lengths, pool_strides):
             out = math.floor(out + 2 * (k // 2) - k + 1)
             out = math.floor((out - pl) / ps + 1)
+            # Ensure output is at least 1 to avoid zero-sized tensors
+            out = max(1, out)
             seq_lengths.append(int(out))
         return seq_lengths
 
@@ -372,6 +503,7 @@ class _ChannelAttentionBlock(nn.Module):
     ----------
     attention_mode : str, optional
         The type of attention mechanism to apply. If `None`, no attention is applied.
+
         - "se" for Squeeze-and-excitation network
         - "gsop" for Global Second-Order Pooling
         - "fca" for Frequency Channel Attention Network

braindecode/models/{sleep_stager_eldele_2021.py → attn_sleep.py}

@@ -8,18 +8,19 @@ from copy import deepcopy
 
 import torch
 import torch.nn.functional as F
+from mne.utils import deprecated
 from torch import nn
 
 from braindecode.models.base import EEGModuleMixin
 from braindecode.modules import CausalConv1d
 
 
-class SleepStagerEldele2021(EEGModuleMixin, nn.Module):
+class AttnSleep(EEGModuleMixin, nn.Module):
     """Sleep Staging Architecture from Eldele et al. (2021) [Eldele2021]_.
 
     .. figure:: https://raw.githubusercontent.com/emadeldeen24/AttnSleep/refs/heads/main/imgs/AttnSleep.png
         :align: center
-        :alt: SleepStagerEldele2021 Architecture
+        :alt: AttnSleep Architecture
 
     Attention based Neural Net for sleep staging as described in [Eldele2021]_.
     The code for the paper and this model is also available at [1]_.
@@ -533,3 +534,12 @@ class _PositionwiseFeedForward(nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Implements FFN equation."""
         return self.w_2(self.dropout(self.activate(self.w_1(x))))
+
+
+@deprecated(
+    "`SleepStagerEldele2021` was renamed to `AttnSleep` in v1.12 to follow original author's name; this alias will be removed in v1.14."
+)
+class SleepStagerEldele2021(AttnSleep):
+    """Deprecated alias for SleepStagerEldele2021."""
+
+    pass

braindecode/models/base.py

@@ -5,15 +5,35 @@
 
 from __future__ import annotations
 
+import json
 import warnings
 from collections import OrderedDict
-from typing import Dict, Iterable, Optional
+from pathlib import Path
+from typing import Dict, Iterable, Optional, Type, Union
 
 import numpy as np
 import torch
 from docstring_inheritance import NumpyDocstringInheritanceInitMeta
+from mne.utils import _soft_import
 from torchinfo import ModelStatistics, summary
 
+from braindecode.version import __version__
+
+huggingface_hub = _soft_import(
+    "huggingface_hub", "Hugging Face Hub integration", strict=False
+)
+
+HAS_HF_HUB = huggingface_hub is not False
+
+
+class _BaseHubMixin:
+    pass
+
+
+# Define base class for hub mixin
+if HAS_HF_HUB:
+    _BaseHubMixin: Type = huggingface_hub.PyTorchModelHubMixin  # type: ignore
+
 
 def deprecated_args(obj, *old_new_args):
     out_args = []
@@ -32,10 +52,14 @@ def deprecated_args(obj, *old_new_args):
     return out_args
 
 
-class EEGModuleMixin(metaclass=NumpyDocstringInheritanceInitMeta):
+class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta):
     """
     Mixin class for all EEG models in braindecode.
 
+    This class integrates with Hugging Face Hub when the ``huggingface_hub`` package
+    is installed, enabling models to be pushed to and loaded from the Hub using
+    :func:`push_to_hub()` and :func:`from_pretrained()` methods.
+
     Parameters
     ----------
     n_outputs : int
@@ -62,8 +86,87 @@ class EEGModuleMixin(metaclass=NumpyDocstringInheritanceInitMeta):
     -----
     If some input signal-related parameters are not specified,
     there will be an attempt to infer them from the other parameters.
+
+    .. rubric:: Hugging Face Hub integration
+
+    When the optional ``huggingface_hub`` package is installed, all models
+    automatically gain the ability to be pushed to and loaded from the
+    Hugging Face Hub. Install with::
+
+        pip install braindecode[hug]
+
+    **Pushing a model to the Hub:**
+
+    .. code-block:: python
+
+        from braindecode.models import EEGNetv4
+
+        # Train your model
+        model = EEGNetv4(n_chans=22, n_outputs=4, n_times=1000)
+        # ... training code ...
+
+        # Push to the Hub
+        model.push_to_hub(
+            repo_id="username/my-eegnet-model", commit_message="Initial model upload"
+        )
+
+    **Loading a model from the Hub:**
+
+    .. code-block:: python
+
+        from braindecode.models import EEGNetv4
+
+        # Load pretrained model
+        model = EEGNetv4.from_pretrained("username/my-eegnet-model")
+
+    The integration automatically handles EEG-specific parameters (n_chans,
+    n_times, sfreq, chs_info, etc.) by saving them in a config file alongside
+    the model weights. This ensures that loaded models are correctly configured
+    for their original data specifications.
+
+    .. important::
+        Currently, only EEG-specific parameters (n_outputs, n_chans, n_times,
+        input_window_seconds, sfreq, chs_info) are saved to the Hub. Model-specific
+        parameters (e.g., dropout rates, activation functions, number of filters)
+        are not preserved and will use their default values when loading from the Hub.
+
+        To use non-default model parameters, specify them explicitly when calling
+        :func:`from_pretrained()`::
+
+            model = EEGNet.from_pretrained("user/model", dropout=0.3, activation='relu')
+
+        Full parameter serialization will be addressed in a future update.
     """
 
+    def __init_subclass__(cls, **kwargs):
+        if not HAS_HF_HUB:
+            super().__init_subclass__(**kwargs)
+            return
+
+        base_tags = ["braindecode", cls.__name__]
+        user_tags = kwargs.pop("tags", None)
+        tags = list(user_tags) if user_tags is not None else []
+        for tag in base_tags:
+            if tag not in tags:
+                tags.append(tag)
+
+        docs_url = kwargs.pop(
+            "docs_url",
+            f"https://braindecode.org/stable/generated/braindecode.models.{cls.__name__}.html",
+        )
+        repo_url = kwargs.pop("repo_url", "https://braindecode.org")
+        library_name = kwargs.pop("library_name", "braindecode")
+        license = kwargs.pop("license", "bsd-3-clause")
+        # TODO: model_card_template can be added in the future for custom model cards
+        super().__init_subclass__(
+            tags=tags,
+            docs_url=docs_url,
+            repo_url=repo_url,
+            library_name=library_name,
+            license=license,
+            **kwargs,
+        )
+
     def __init__(
         self,
         n_outputs: Optional[int] = None,  # type: ignore[assignment]
@@ -73,6 +176,16 @@ class EEGModuleMixin(metaclass=NumpyDocstringInheritanceInitMeta):
         input_window_seconds: Optional[float] = None,  # type: ignore[assignment]
         sfreq: Optional[float] = None,  # type: ignore[assignment]
     ):
+        # Deserialize chs_info if it comes as a list of dicts (from Hub)
+        if chs_info is not None and isinstance(chs_info, list):
+            if len(chs_info) > 0 and isinstance(chs_info[0], dict):
+                # Check if it needs deserialization (has 'loc' as list)
+                if "loc" in chs_info[0] and isinstance(chs_info[0]["loc"], list):
+                    chs_info = self._deserialize_chs_info(chs_info)
+                    warnings.warn(
+                        "Modifying chs_info argument using the _deserialize_chs_info() method"
+                    )
+
         if n_chans is not None and chs_info is not None and len(chs_info) != n_chans:
             raise ValueError(f"{n_chans=} different from {chs_info=} length")
         if (
@@ -294,3 +407,168 @@ class EEGModuleMixin(metaclass=NumpyDocstringInheritanceInitMeta):
 
     def __str__(self) -> str:
         return str(self.get_torchinfo_statistics())
+
+    @staticmethod
+    def _serialize_chs_info(chs_info):
+        """
+        Serialize MNE channel info to JSON-compatible format.
+
+        Parameters
+        ----------
+        chs_info : list of dict or None
+            Channel information from MNE Info object.
+
+        Returns
+        -------
+        list of dict or None
+            Serialized channel information that can be saved to JSON.
+        """
+        if chs_info is None:
+            return None
+
+        serialized = []
+        for ch in chs_info:
+            # Extract serializable fields from MNE channel info
+            ch_dict = {
+                "ch_name": ch.get("ch_name", ""),
+            }
+
+            # Handle kind field - can be either string or integer
+            kind_val = ch.get("kind")
+            if kind_val is not None:
+                ch_dict["kind"] = (
+                    kind_val if isinstance(kind_val, str) else int(kind_val)
+                )
+
+            # Add numeric fields with safe conversion
+            coil_type = ch.get("coil_type")
+            if coil_type is not None:
+                ch_dict["coil_type"] = int(coil_type)
+
+            unit = ch.get("unit")
+            if unit is not None:
+                ch_dict["unit"] = int(unit)
+
+            cal = ch.get("cal")
+            if cal is not None:
+                ch_dict["cal"] = float(cal)
+
+            range_val = ch.get("range")
+            if range_val is not None:
+                ch_dict["range"] = float(range_val)
+
+            # Serialize location array if present
+            if "loc" in ch and ch["loc"] is not None:
+                ch_dict["loc"] = (
+                    ch["loc"].tolist()
+                    if hasattr(ch["loc"], "tolist")
+                    else list(ch["loc"])
+                )
+            serialized.append(ch_dict)
+
+        return serialized
+
+    @staticmethod
+    def _deserialize_chs_info(chs_info_dict):
+        """
+        Deserialize channel info from JSON-compatible format to MNE-like structure.
+
+        Parameters
+        ----------
+        chs_info_dict : list of dict or None
+            Serialized channel information.
+
+        Returns
+        -------
+        list of dict or None
+            Deserialized channel information compatible with MNE.
+        """
+        if chs_info_dict is None:
+            return None
+
+        deserialized = []
+        for ch_dict in chs_info_dict:
+            ch = ch_dict.copy()
+            # Convert location back to numpy array if present
+            if "loc" in ch and ch["loc"] is not None:
+                ch["loc"] = np.array(ch["loc"])
+            deserialized.append(ch)
+
+        return deserialized
+
+    def _save_pretrained(self, save_directory):
+        """
+        Save model configuration and weights to the Hub.
+
+        This method is called by PyTorchModelHubMixin.push_to_hub() to save
+        model-specific configuration alongside the model weights.
+
+        Parameters
+        ----------
+        save_directory : str or Path
+            Directory where the configuration should be saved.
+        """
+        if not HAS_HF_HUB:
+            return
+
+        save_directory = Path(save_directory)
+
+        # Collect EEG-specific configuration
+        config = {
+            "n_outputs": self._n_outputs,
+            "n_chans": self._n_chans,
+            "n_times": self._n_times,
+            "input_window_seconds": self._input_window_seconds,
+            "sfreq": self._sfreq,
+            "chs_info": self._serialize_chs_info(self._chs_info),
+            "braindecode_version": __version__,
+        }
+
+        # Save to config.json
+        config_path = save_directory / "config.json"
+        with open(config_path, "w") as f:
+            json.dump(config, f, indent=2)
+
+        # Save model weights with standard Hub filename
+        weights_path = save_directory / "pytorch_model.bin"
+        torch.save(self.state_dict(), weights_path)
+
+        # Also save in safetensors format using parent's implementation
+        try:
+            super()._save_pretrained(save_directory)
+        except (ImportError, RuntimeError) as e:
+            # Fallback to pytorch_model.bin if safetensors saving fails
+            warnings.warn(
+                f"Could not save model in safetensors format: {e}. "
+                "Model weights saved in pytorch_model.bin instead.",
+                stacklevel=2,
+            )
+
+    if HAS_HF_HUB:
+
+        @classmethod
+        def _from_pretrained(
+            cls,
+            *,
+            model_id: str,
+            revision: Optional[str],
+            cache_dir: Optional[Union[str, Path]],
+            force_download: bool,
+            local_files_only: bool,
+            token: Union[str, bool, None],
+            map_location: str = "cpu",
+            strict: bool = False,
+            **model_kwargs,
+        ):
+            model_kwargs.pop("braindecode_version", None)
+            return super()._from_pretrained(  # type: ignore
+                model_id=model_id,
+                revision=revision,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                local_files_only=local_files_only,
+                token=token,
+                map_location=map_location,
+                strict=strict,
+                **model_kwargs,
+            )