braindecode 1.6.1.dev1052__tar.gz → 1.6.1.dev1053__tar.gz

{braindecode-1.6.1.dev1052/braindecode.egg-info → braindecode-1.6.1.dev1053}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.6.1.dev1052
+Version: 1.6.1.dev1053
 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.dev1052 → braindecode-1.6.1.dev1053}/braindecode/eegneuralnet.py

@@ -7,6 +7,7 @@
 import abc
 import inspect
 import logging
+import warnings
 from typing import Literal
 
 import mne
@@ -41,6 +42,7 @@ def _get_model(model: str):
 
 class _EEGNeuralNet(NeuralNet, abc.ABC):
     signal_args_set_ = False
+    _warned_drop_last_ = False
 
     @property
     def log(self):
@@ -267,6 +269,44 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
             X = X.get_data(units="uV")
         return super().get_dataset(X, y)
 
+    def get_iterator(self, dataset, training=False):
+        iterator = super().get_iterator(dataset, training=training)
+        # Warn (once per fit) if drop_last would discard every training batch,
+        # e.g. a training set smaller than batch_size. This silently skips
+        # training without raising any error, which is hard to diagnose.
+        if training and not self._warned_drop_last_:
+            self._warn_if_training_batches_dropped(iterator)
+        return iterator
+
+    def _warn_if_training_batches_dropped(self, loader):
+        """Warn if ``drop_last`` discards all training batches.
+
+        With ``iterator_train__drop_last=True`` (the braindecode default), a
+        training set smaller than ``batch_size`` yields zero batches, so no
+        forward/backward pass runs and the model is left untrained without any
+        error. We inspect the resolved training ``DataLoader`` (after
+        ``train_split`` has been applied), so the count reflects the actual
+        number of training examples.
+        """
+        self._warned_drop_last_ = True
+        if not getattr(loader, "drop_last", False):
+            return
+        batch_size = getattr(loader, "batch_size", None)
+        try:
+            n_train = len(loader.dataset)
+        except TypeError:
+            return  # e.g. an IterableDataset whose length is unknown
+        if not batch_size or n_train == 0 or n_train >= batch_size:
+            return
+        warnings.warn(
+            f"The training set has {n_train} example(s), which is smaller than "
+            f"batch_size={batch_size}. With iterator_train__drop_last=True "
+            f"(the braindecode default), all training batches are dropped and "
+            f"the model is not trained. Reduce batch_size to at most {n_train}, "
+            f"or pass iterator_train__drop_last=False.",
+            UserWarning,
+        )
+
     def partial_fit(self, X, y=None, classes=None, **fit_params):
         """Fit the module.
 
@@ -366,6 +406,7 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
           the module and to the ``self.train_split`` call.
         """
         # this needs to be executed before the net is initialized:
+        self._warned_drop_last_ = False  # re-check on every fresh fit
         if not self.signal_args_set_:
             self._set_signal_args(X, y, classes=None)
             self.signal_args_set_ = True

{braindecode-1.6.1.dev1052 → braindecode-1.6.1.dev1053}/braindecode/models/__init__.py

@@ -15,6 +15,7 @@ from .deep4 import Deep4Net
 from .deepsleepnet import DeepSleepNet
 from .dgcnn import DGCNN
 from .eegconformer import EEGConformer
+from .eegdino import EEGDINO
 from .eeginception_erp import EEGInceptionERP
 from .eeginception_mi import EEGInceptionMI
 from .eegitnet import EEGITNet
@@ -84,6 +85,7 @@ __all__ = [
     "DeepSleepNet",
     "BrainModule",
     "EEGConformer",
+    "EEGDINO",
     "EEGPT",
     "InterpolatedEEGPT",
     "EEGInceptionERP",

braindecode-1.6.1.dev1053/braindecode/models/eegdino.py

@@ -0,0 +1,533 @@
+# Authors: Bruno Aristimunha <b.aristimunha@gmail.com>
+#
+# Code adapted from https://huggingface.co/eegdino/EEG-DINO
+#
+# License: BSD (3-clause)
+from __future__ import annotations
+
+from typing import Sequence
+from warnings import warn
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from braindecode.models.base import EEGModuleMixin
+from braindecode.modules import DropPath, PatchTokenizer
+
+
+class EEGDINO(EEGModuleMixin, nn.Module):
+    r"""EEG-DINO from Wang et al. (2025) [eegdino]_.
+
+    :bdg-danger:`Foundation Model` :bdg-info:`Attention/Transformer`
+
+    .. figure:: https://braindecode.org/dev/_static/model/eegdino.png
+        :align: center
+        :alt: EEG-DINO Architecture
+
+    EEG-DINO is a ViT-style EEG foundation model pre-trained with DINO-v2
+    hierarchical self-distillation. Only the *encoder* (plus a classification
+    head) is integrated here; the self-distillation pre-training is out of scope
+    and it was not released by the authors.
+
+    The forward path is, end to end:
+
+    ``(batch, n_chans, n_times)`` → patchify → time-frequency embedding →
+    decoupled positional embedding → transformer encoder (+ global token) →
+    pooling head → ``(batch, n_outputs)``.
+
+    .. rubric:: Step 1 -- Patchify
+
+    The signal is split along time into non-overlapping patches of
+    ``patch_size`` samples (200 samples = 1 second at 200 Hz), giving one
+    *token* per (channel, patch). Inputs whose length is not a multiple of
+    ``patch_size`` are zero-padded with a warning.
+    This is analogous to the patchification in ViT, BEiT (LaBraM), and CBraMod.
+
+    .. rubric:: Step 2 -- Time-Frequency Embedding (TFE)
+
+    Each patch is embedded by summing two branches, exactly as in
+    :class:`~braindecode.models.CBraMod` (the EEG-DINO paper reuses CBraMod's
+    TFE): a **time-domain** branch of stacked grouped convolutions
+    (``proj_in``), and a **frequency-domain** branch projecting the magnitude
+    of the patch's real FFT (``spectral_proj``). The embedding dimension
+    ``emb_dim`` is therefore *derived* from the convolution configuration, not
+    set independently.
+
+    .. rubric:: Step 3 -- Decoupled Positional Embedding (DPE)
+
+    Where CBraMod uses a single convolutional positional encoding (ACPE),
+    EEG-DINO *decouples* space and time and adds both to every token: a
+    learnable **one-hot channel embedding** over a fixed vocabulary of
+    ``n_channel_embeddings`` slots (input channel ``i`` maps to slot ``i``, so a
+    19-slot embedding serves any ``n_chans <= 19``) and a **depthwise temporal
+    convolution** over the patch axis (``time_encoding``).
+
+    .. rubric:: Step 4 -- Transformer Encoder & Global Token
+
+    Tokens are flattened to a single sequence and processed by ``n_layer``
+    pre-norm transformer blocks (BEiT-style attention with separate query/value
+    biases). A learnable ``global_tokens`` summary token is prepended after the
+    ``global_token_layer``-th block and attends jointly with the patch tokens.
+
+    .. rubric:: Step 5 -- Classification Head (``final_layer``)
+
+    Following EEG-DINO's finetuning head, the patch tokens (global token
+    excluded) pass through a per-token projection, are mean-pooled over channels
+    then over patches (with an intermediate projection), and a three-layer MLP
+    (``final_layer``) maps them to ``n_outputs``
+    (:class:`~braindecode.EEGClassifier` applies the softmax). With
+    ``return_encoder_output=True`` the mean-pooled encoder representation is
+    returned instead (linear probing).
+
+    .. warning::
+       EEG-DINO was pre-trained on EEG in microvolts scaled by ``1 / 100``. The
+       model does **not** rescale its input, so for the released weights to
+       behave as intended provide data on the same scale (e.g. microvolts
+       divided by 100) or normalize your signals consistently.
+
+    .. important::
+       **Pre-trained Weights Available**
+
+       Small and Medium encoders converted from the released checkpoints are
+       hosted on the Hugging Face Hub, one repository per size. Only the encoder
+       is pretrained; the classification head is randomly initialized, so
+       fine-tune or linear-probe before use::
+
+           from braindecode.models import EEGDINO
+
+           model = EEGDINO.from_pretrained(
+               "braindecode/eegdino-small-pretrained",  # or -medium-pretrained
+               n_outputs=6,
+               n_chans=19,
+               sfreq=200,
+           )
+
+       The Small/Medium/Large architectures are also available in
+       :data:`EEGDINO_CONFIGS`. Requires ``braindecode[hub]``.
+
+    .. versionadded:: 1.6.1
+
+    Parameters
+    ----------
+    patch_size : int, default=200
+        Temporal patch size in samples (200 = 1 second at 200 Hz). Fixed at 200
+        for the released weights (the FFT branch uses ``patch_size // 2 + 1`` bins).
+    n_layer : int, default=12
+        Number of transformer encoder layers.
+    nhead : int, default=8
+        Number of attention heads.
+    dim_feedforward : int, default=512
+        Hidden size of the transformer feed-forward block.
+    channels_kernel_stride_padding_norm : sequence of tuple, optional
+        Configuration of the time-domain convolutions in the patch embedding,
+        as ``(out_channels, kernel, stride, padding, (groups, group_channels))``
+        per layer. The embedding dimension is derived from this (see
+        :class:`~braindecode.models.CBraMod`). Default is the EEG-DINO-Small /
+        CBraMod configuration.
+    n_channel_embeddings : int, default=19
+        Size of the one-hot channel-embedding vocabulary. ``n_chans`` must not
+        exceed it; the first ``n_chans`` slots are used. Default 19 (the released
+        montage), so the pretrained weights load for any ``n_chans <= 19``.
+    n_global_tokens : int, default=1
+        Number of learnable global summary tokens.
+    global_token_layer : int, default=1
+        1-based index of the encoder layer after which the global tokens are
+        inserted.
+    activation : type[nn.Module], default=nn.GELU
+        Activation function used in the transformer feed-forward blocks and the
+        classification head.
+    patch_activation : type[nn.Module], default=nn.GELU
+        Activation function used in the patch-embedding convolutions. Defaults to
+        ``nn.GELU`` to match the released pretrained weights.
+    drop_prob : float, default=0.1
+        Dropout / stochastic-depth probability in the encoder.
+    return_features : bool, default=False
+        If True, ``forward`` returns ``{"features", "cls_token"}``.
+    return_encoder_output : bool, default=False
+        If True, ``final_layer`` is :class:`~torch.nn.Identity` and ``forward``
+        returns the pooled encoder representation (linear probing).
+
+    References
+    ----------
+    .. [eegdino] Wang, X., Liu, X., Liu, X., Si, Q., Xu, Z., Li, Y., & Zhen, X.
+       (2025). EEG-DINO: Learning EEG Foundation Models via Hierarchical
+       Self-Distillation. In Medical Image Computing and Computer Assisted
+       Intervention (MICCAI 2025).
+    """
+
+    def __init__(
+        self,
+        n_outputs=None,
+        n_chans=None,
+        chs_info=None,
+        n_times=None,
+        input_window_seconds=None,
+        sfreq=None,
+        patch_size: int = 200,
+        n_layer: int = 12,
+        nhead: int = 8,
+        dim_feedforward: int = 512,
+        channels_kernel_stride_padding_norm: Sequence[
+            tuple[int, int, int, int, tuple[int, int]]
+        ] = (
+            (25, 49, 25, 24, (5, 25)),
+            (25, 3, 1, 1, (5, 25)),
+            (25, 3, 1, 1, (5, 25)),
+        ),
+        n_channel_embeddings: int = 19,
+        n_global_tokens: int = 1,
+        global_token_layer: int = 1,
+        activation: type[nn.Module] = nn.GELU,
+        patch_activation: type[nn.Module] = nn.GELU,
+        drop_prob: float = 0.1,
+        return_features: bool = False,
+        return_encoder_output: bool = False,
+    ):
+        super().__init__(
+            n_outputs=n_outputs,
+            n_chans=n_chans,
+            chs_info=chs_info,
+            n_times=n_times,
+            input_window_seconds=input_window_seconds,
+            sfreq=sfreq,
+        )
+        del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
+
+        if not 1 <= global_token_layer <= n_layer:
+            raise ValueError(
+                f"global_token_layer must be in [1, n_layer={n_layer}], "
+                f"got {global_token_layer}."
+            )
+        if n_global_tokens < 1:
+            raise ValueError(f"n_global_tokens must be >= 1, got {n_global_tokens}.")
+        if self.n_chans > n_channel_embeddings:
+            raise ValueError(
+                f"n_chans ({self.n_chans}) must not exceed n_channel_embeddings "
+                f"({n_channel_embeddings}); the released weights use 19."
+            )
+
+        if self._sfreq is not None and self.sfreq != 200:
+            warn(
+                f"EEG-DINO was trained at 200 Hz but sfreq={self.sfreq}. Inputs are "
+                "not resampled internally; results may be unreliable.",
+                UserWarning,
+            )
+
+        self.tokenizer = PatchTokenizer(
+            patch_size, n_times=self.n_times, learnable=False
+        )
+        self.n_global_tokens = n_global_tokens
+        self.global_token_layer = global_token_layer
+        self.activation = activation
+        self.return_features = return_features
+        self.return_encoder_output = return_encoder_output
+
+        self.patch_embedding = _PatchEmbedding(
+            patch_size=patch_size,
+            channels_kernel_stride_padding_norm=channels_kernel_stride_padding_norm,
+            n_channel_embeddings=n_channel_embeddings,
+            activation=patch_activation,
+            drop_prob=drop_prob,
+        )
+        self.emb_dim = self.patch_embedding.emb_dim
+
+        self.encoder_layers = nn.ModuleList(
+            _TransformerEncoderLayer(
+                self.emb_dim, nhead, dim_feedforward, activation, drop_prob
+            )
+            for _ in range(n_layer)
+        )
+        self.global_tokens = nn.Parameter(torch.zeros(1, n_global_tokens, self.emb_dim))
+
+        self.final_layer = self._make_head()
+
+        self._init_weights()
+
+    def _init_weights(self):
+        nn.init.trunc_normal_(self.global_tokens, std=0.02)
+
+    def _make_head(self):
+        if self.return_encoder_output:
+            return nn.Identity()  # ``n_outputs`` not needed for linear probing
+        return _ClassificationHead(self.emb_dim, self.n_outputs, self.activation)
+
+    def reset_head(self, n_outputs):
+        """Replace ``final_layer`` with a fresh classification head for ``n_outputs``.
+
+        Asking for a head implies a classification model, so this also clears
+        ``return_encoder_output`` (a linear-probe model thereby gains a real head),
+        mirroring :meth:`braindecode.models.CBraMod.reset_head`.
+        """
+        self._n_outputs = n_outputs
+        self.return_encoder_output = False
+        self.final_layer = self._make_head()
+
+    def forward(self, x, return_features: bool | None = None):
+        """Forward pass.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input of shape ``(batch, n_chans, n_times)``. The model does not
+            rescale ``x`` (see the amplitude-scale warning in the class docstring).
+        return_features : bool, optional
+            Overrides ``self.return_features`` for this call.
+
+        Returns
+        -------
+        torch.Tensor or dict
+            Logits ``(batch, n_outputs)``; or, with features,
+            ``{"features": (batch, n_chans * n_patches, emb_dim),
+            "cls_token": (batch, emb_dim)}``.
+        """
+        if return_features is None:
+            return_features = self.return_features
+        x = self.tokenizer(x)  # (batch, n_chans, n_patches, patch_size)
+
+        patch_emb = self.patch_embedding(x)  # (batch, n_chans, n_patches, emb_dim)
+        n_chans = patch_emb.shape[1]
+        tokens = patch_emb.flatten(1, 2)  # (batch, chans * patches, emb)
+
+        global_tokens = self.global_tokens.expand(tokens.shape[0], -1, -1)
+        for i, layer in enumerate(self.encoder_layers):
+            tokens = layer(tokens)
+            if i + 1 == self.global_token_layer:
+                tokens = torch.cat([global_tokens, tokens], dim=1)
+
+        global_out = tokens[:, : self.n_global_tokens]
+        patch_out = tokens[:, self.n_global_tokens :]
+
+        if return_features:
+            return {"features": patch_out, "cls_token": global_out.mean(dim=1)}
+
+        if self.return_encoder_output:
+            return patch_out.mean(dim=1)  # pooled encoder features (linear probing)
+        return self.final_layer(patch_out, n_chans)
+
+
+#: Architecture presets. Small and Medium have released weights; Large does not.
+EEGDINO_CONFIGS = {
+    "small": dict(
+        n_layer=12,
+        nhead=8,
+        dim_feedforward=512,
+        channels_kernel_stride_padding_norm=(
+            (25, 49, 25, 24, (5, 25)),
+            (25, 3, 1, 1, (5, 25)),
+            (25, 3, 1, 1, (5, 25)),
+        ),
+    ),
+    "medium": dict(
+        n_layer=16,
+        nhead=8,
+        dim_feedforward=1024,
+        channels_kernel_stride_padding_norm=(
+            (64, 49, 25, 24, (8, 64)),
+            (128, 3, 1, 1, (8, 128)),
+            (64, 3, 1, 1, (8, 64)),
+        ),
+    ),
+    "large": dict(
+        n_layer=24,
+        nhead=8,
+        dim_feedforward=2048,
+        channels_kernel_stride_padding_norm=(
+            (128, 49, 25, 24, (16, 128)),
+            (256, 3, 1, 1, (16, 256)),
+            (128, 3, 1, 1, (16, 128)),
+        ),
+    ),
+}
+
+
+class _PatchEmbedding(nn.Module):
+    """Time-frequency patch embedding with a decoupled positional embedding.
+
+    The time-domain (``proj_in``) and frequency-domain (``spectral_proj``)
+    branches are CBraMod's TFE (see :class:`~braindecode.models.CBraMod`); the
+    one-hot ``channel_embedding`` and depthwise ``time_encoding`` are EEG-DINO's
+    decoupled positional embedding.
+    """
+
+    def __init__(
+        self,
+        patch_size,
+        channels_kernel_stride_padding_norm,
+        n_channel_embeddings,
+        activation=nn.GELU,
+        drop_prob=0.1,
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        self.channels_kernel_stride_padding_norm = channels_kernel_stride_padding_norm
+        self.n_channel_embeddings = n_channel_embeddings
+
+        in_channels, conv_layers = 1, []
+        for (
+            out_channels,
+            kernel,
+            stride,
+            padding,
+            norm,
+        ) in channels_kernel_stride_padding_norm:
+            conv_layers += [
+                nn.Conv2d(
+                    in_channels, out_channels, (1, kernel), (1, stride), (0, padding)
+                ),
+                nn.GroupNorm(*norm),
+                activation(),
+            ]
+            in_channels = out_channels
+        self.proj_in = nn.Sequential(*conv_layers)
+
+        self.spectral_proj = nn.Sequential(
+            nn.Linear(patch_size // 2 + 1, self.emb_dim), nn.Dropout(drop_prob)
+        )
+        self.channel_embedding = nn.Linear(n_channel_embeddings, self.emb_dim)
+        # Sequential wrapper keeps the released-checkpoint key ``time_encoding.0.*``.
+        self.time_encoding = nn.Sequential(
+            nn.Conv2d(
+                self.emb_dim, self.emb_dim, (1, 5), (1, 1), (0, 2), groups=self.emb_dim
+            )
+        )
+
+    @property
+    def emb_dim(self):
+        """Embedding dimension implied by the convolution configuration."""
+        reduced = self.patch_size
+        for _, kernel, stride, padding, _ in self.channels_kernel_stride_padding_norm:
+            reduced = (reduced + 2 * padding - kernel) // stride + 1
+        return self.channels_kernel_stride_padding_norm[-1][0] * reduced
+
+    def forward(self, x):
+        # x: (batch, n_chans, n_patches, patch_size)
+        n_chans = x.shape[1]
+
+        # Time-frequency embedding (CBraMod TFE): conv branch + FFT-magnitude branch.
+        time_tokens = self.proj_in(x.flatten(1, 2).unsqueeze(1))
+        batch, conv_dim, _, conv_len = time_tokens.shape
+        time_tokens = (
+            time_tokens.reshape(batch, conv_dim, n_chans, -1, conv_len)
+            .permute(0, 2, 3, 1, 4)
+            .flatten(3, 4)
+        )
+        spectrum = torch.fft.rfft(x, dim=-1, norm="forward").abs()
+        patch_emb = time_tokens + self.spectral_proj(spectrum)
+
+        # Decoupled positional embedding: one-hot channel + depthwise temporal conv.
+        # The one-hot uses the first ``n_chans`` of the ``n_channel_embeddings``
+        # slots, so the released 19-slot embedding serves any n_chans <= 19.
+        channel_ids = torch.arange(n_chans, device=x.device)
+        one_hot = F.one_hot(channel_ids, self.n_channel_embeddings).to(
+            self.channel_embedding.weight.dtype
+        )
+        channel_emb = self.channel_embedding(one_hot)
+        patch_emb = patch_emb + channel_emb.unsqueeze(0).unsqueeze(2)
+        temporal_emb = self.time_encoding(patch_emb.permute(0, 3, 1, 2))
+        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)
+        self.attn = _Attention(emb_dim, nhead, attn_drop=drop_prob, proj_drop=drop_prob)
+        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 _Attention(nn.Module):
+    """BEiT-style attention: fused ``qkv`` without bias plus separate q/v biases.
+
+    ``head_dim`` is decoupled from ``emb_dim`` (``all_head_dim = head_dim *
+    nhead``) so a custom ``nhead`` need not divide ``emb_dim``. For the released
+    presets ``nhead=8`` divides every ``emb_dim`` (200/512/1024), so this is a
+    no-op there and the released ``qkv``/``proj`` shapes are preserved.
+    """
+
+    def __init__(self, emb_dim, nhead=8, attn_drop=0.0, proj_drop=0.0):
+        super().__init__()
+        if not 1 <= nhead <= emb_dim:
+            raise ValueError(
+                f"nhead ({nhead}) must be in [1, emb_dim={emb_dim}] so head_dim >= 1."
+            )
+        self.nhead = nhead
+        head_dim = emb_dim // nhead
+        all_head_dim = head_dim * nhead
+        self.scale = head_dim**-0.5
+        self.qkv = nn.Linear(emb_dim, all_head_dim * 3, bias=False)
+        self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
+        self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(all_head_dim, emb_dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias), self.v_bias))
+        qkv = F.linear(x, self.qkv.weight, bias)
+        batch, tokens, _ = qkv.shape
+        query, key, value = (
+            qkv.reshape(batch, tokens, 3, self.nhead, -1)
+            .permute(2, 0, 3, 1, 4)
+            .unbind(0)
+        )
+        attn = (query * self.scale @ key.transpose(-2, -1)).softmax(dim=-1)
+        attn = self.attn_drop(attn)
+        out = (attn @ value).transpose(1, 2).flatten(2)
+        return self.proj_drop(self.proj(out))
+
+
+class _MLP(nn.Module):
+    """Feed-forward block; ``fc1``/``fc2`` names match the released checkpoint."""
+
+    def __init__(self, emb_dim, hidden_dim, activation=nn.GELU, drop=0.0):
+        super().__init__()
+        self.fc1 = nn.Linear(emb_dim, hidden_dim)
+        self.act = activation()
+        self.fc2 = nn.Linear(hidden_dim, emb_dim)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        return self.drop(self.fc2(self.act(self.fc1(x))))
+
+
+class _ClassificationHead(nn.Module):
+    """EEG-DINO finetuning head (Wang et al., 2025).
+
+    Applies a per-token projection, mean-pools the patch tokens over channels
+    then over patches (with an intermediate projection), and maps the pooled
+    ``emb_dim`` vector to ``n_outputs`` with a three-layer MLP.
+    """
+
+    def __init__(
+        self, emb_dim, n_outputs, activation=nn.GELU, drop_prob_1=0.5, drop_prob_2=0.3
+    ):
+        super().__init__()
+        self.token_proj = nn.Sequential(nn.Linear(emb_dim, emb_dim), activation())
+        self.time_proj = nn.Sequential(nn.Linear(emb_dim, emb_dim), activation())
+        self.classifier = nn.Sequential(
+            nn.Linear(emb_dim, emb_dim // 2),
+            activation(),
+            nn.Dropout(drop_prob_1),
+            nn.Linear(emb_dim // 2, emb_dim // 4),
+            activation(),
+            nn.Dropout(drop_prob_2),
+            nn.Linear(emb_dim // 4, n_outputs),
+        )
+
+    def forward(self, patch_tokens, n_chans):
+        x = self.token_proj(patch_tokens)
+        x = x.reshape(x.shape[0], n_chans, -1, x.shape[2]).mean(dim=1)
+        x = self.time_proj(x).mean(dim=1)
+        return self.classifier(x)

{braindecode-1.6.1.dev1052 → braindecode-1.6.1.dev1053}/braindecode/models/summary.csv

@@ -56,3 +56,4 @@ 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"
+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"

{braindecode-1.6.1.dev1052 → braindecode-1.6.1.dev1053}/braindecode/models/util.py

@@ -432,6 +432,7 @@ models_mandatory_parameters: list[
         {"n_chans": 19, "n_times": 6000},
     ),
     ("DGCNN", ["n_chans", "n_outputs", "n_times", "chs_info"], None),
+    ("EEGDINO", ["n_chans", "n_outputs", "n_times"], None),
 ]
 
 ################################################################

{braindecode-1.6.1.dev1052 → braindecode-1.6.1.dev1053}/braindecode/modules/__init__.py

@@ -14,7 +14,7 @@ from .attention import (
     MultiHeadAttention,
     SqueezeAndExcitation,
 )
-from .blocks import MLP, FeedForwardBlock, InceptionBlock
+from .blocks import MLP, FeedForwardBlock, InceptionBlock, PatchTokenizer
 from .convolution import (
     AvgPool2dWithConv,
     CausalConv1d,
@@ -68,6 +68,7 @@ __all__ = [
     "MLP",
     "FeedForwardBlock",
     "InceptionBlock",
+    "PatchTokenizer",
     "AvgPool2dWithConv",
     "CausalConv1d",
     "CombinedConv",