Braindecode 1.3.2.dev182544824__tar.gz → 1.3.2.dev183191339__tar.gz

{braindecode-1.3.2.dev182544824/braindecode.egg-info → braindecode-1.3.2.dev183191339}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.3.2.dev182544824
+Version: 1.3.2.dev183191339
 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.3.2.dev182544824 → braindecode-1.3.2.dev183191339}/braindecode/models/base.py

@@ -59,7 +59,25 @@ def deprecated_args(obj, *old_new_args):
     return out_args
 
 
-class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta):
+class _BraindecodeDocstringMeta(NumpyDocstringInheritanceInitMeta):
+    """Defer ``__init__`` wrapping until after docstring inheritance.
+
+    ``NumpyDocstringInheritanceInitMeta`` uses ``inspect.unwrap()``
+    internally, which bypasses ``@wraps`` wrappers.  By wrapping
+    ``__init__`` *after* docstring processing, the metaclass sees the
+    unwrapped function and correctly inherits ``cls.__doc__``.
+    """
+
+    def __init__(cls, class_name, class_bases, class_dict):
+        super().__init__(class_name, class_bases, class_dict)
+        # Only wrap subclass __init__s, not EEGModuleMixin itself.
+        # Wrapping the mixin would cause super().__init__() calls to
+        # overwrite _braindecode_init_kwargs captured by the subclass.
+        if any(isinstance(b, _BraindecodeDocstringMeta) for b in class_bases):
+            track_model_init_kwargs(cls)
+
+
+class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
     """
     Mixin class for all EEG models in braindecode.
 
@@ -132,20 +150,46 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
         # Load pretrained model
         model = {name}.from_pretrained("username/my-{name_lower}-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.
+        # Load with a different number of outputs (head is rebuilt automatically)
+        model = {name}.from_pretrained("username/my-{name_lower}-model", n_outputs=4)
+
+    **Extracting features and replacing the head:**
+
+    .. code-block::
+
+        import torch
+
+        x = torch.randn(1, model.n_chans, model.n_times)
+        # Extract encoder features (consistent dict across all models)
+        out = model(x, return_features=True)
+        features = out["features"]
+
+        # Replace the classification head
+        model.reset_head(n_outputs=10)
+
+    **Saving and restoring full configuration:**
+
+    .. code-block::
+
+        import json
+
+        config = model.get_config()            # all __init__ params
+        with open("config.json", "w") as f:
+            json.dump(config, f)
+
+        model2 = {name}.from_config(config)    # reconstruct (no weights)
 
     All model parameters (both EEG-specific and model-specific such as
     dropout rates, activation functions, number of filters) are automatically
     saved to the Hub and restored when loading.
+
+    See :ref:`load-pretrained-models` for a complete tutorial.
     """
 
     def __init_subclass__(cls, **kwargs):
         # Append model-specific Hub integration notes to the docstring.
-        # This runs after the metaclass, so we concatenate rather than
-        # override any existing Notes section in the subclass.
+        # This runs before the metaclass __init__, so the Hub notes will
+        # be included in the docstring that the metaclass processes.
         if cls.__doc__ is not None:
             hub_notes = cls._HUB_NOTES_TEMPLATE.format(
                 name=cls.__name__,
@@ -155,7 +199,6 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
 
         if not HAS_HF_HUB:
             super().__init_subclass__(**kwargs)
-            track_model_init_kwargs(cls)
             return
 
         base_tags = ["braindecode", cls.__name__]
@@ -195,7 +238,6 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
             coders=coders,
             **kwargs,
         )
-        track_model_init_kwargs(cls)
 
     def __init__(
         self,
@@ -426,6 +468,34 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
         resolve_type_kwargs(cls, config)
         return cls(**config)
 
+    def reset_head(self, n_outputs):
+        """Replace the classification head for a new number of outputs.
+
+        This is called automatically by :meth:`from_pretrained` when the
+        user passes an ``n_outputs`` that differs from the saved config.
+        Override in subclasses that need a model-specific head structure.
+
+        Parameters
+        ----------
+        n_outputs : int
+            New number of output classes.
+
+        Examples
+        --------
+        >>> from braindecode.models import BENDR
+        >>> model = BENDR(n_chans=22, n_times=1000, n_outputs=4)
+        >>> model.reset_head(10)
+        >>> model.n_outputs
+        10
+
+        .. versionadded:: 1.4
+        """
+        raise NotImplementedError(
+            f"{type(self).__name__} does not implement reset_head(). "
+            "Override this method to support changing n_outputs after "
+            "loading pretrained weights."
+        )
+
     mapping: Optional[Dict[str, str]] = None
 
     def load_state_dict(self, state_dict, *args, **kwargs):
@@ -625,15 +695,78 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
             **model_kwargs,
         ):
             model_kwargs.pop("braindecode_version", None)
+            filename = model_kwargs.pop("filename", None)
             resolve_type_kwargs(cls, model_kwargs)
-            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,
-            )
+
+            # Read saved n_outputs from config.json to detect when the
+            # user wants a different number of outputs.  Works for both
+            # local directories and Hub repo IDs.
+            saved_n_outputs = None
+            try:
+                if Path(model_id).is_dir():
+                    config_file = Path(model_id) / "config.json"
+                else:
+                    config_file = huggingface_hub.hf_hub_download(
+                        repo_id=model_id,
+                        filename="config.json",
+                        revision=revision,
+                        cache_dir=cache_dir,
+                        force_download=force_download,
+                        token=token,
+                        local_files_only=local_files_only,
+                    )
+                with open(config_file, "r") as f:
+                    saved_n_outputs = json.load(f).get("n_outputs")
+            except (OSError, json.JSONDecodeError, KeyError):
+                pass  # config unavailable; skip reset_head logic
+
+            requested_n_outputs = model_kwargs.get("n_outputs")
+
+            # If the user requests different n_outputs, load with the
+            # saved value first (so weights match), then swap the head.
+            if (
+                saved_n_outputs is not None
+                and requested_n_outputs is not None
+                and requested_n_outputs != saved_n_outputs
+            ):
+                model_kwargs["n_outputs"] = saved_n_outputs
+
+            # If a custom filename is provided, temporarily override the
+            # HuggingFace constant so the parent class downloads the
+            # correct file (e.g. "LUNA_base.safetensors" instead of
+            # "model.safetensors").
+            hf_constants = huggingface_hub.constants
+            _orig_safetensors = hf_constants.SAFETENSORS_SINGLE_FILE
+            if filename is not None:
+                hf_constants.SAFETENSORS_SINGLE_FILE = filename
+            try:
+                model = 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,
+                )
+            finally:
+                hf_constants.SAFETENSORS_SINGLE_FILE = _orig_safetensors
+
+            if (
+                saved_n_outputs is not None
+                and requested_n_outputs is not None
+                and requested_n_outputs != saved_n_outputs
+            ):
+                try:
+                    model.reset_head(requested_n_outputs)
+                except NotImplementedError:
+                    raise ValueError(
+                        f"{type(model).__name__} does not support changing "
+                        f"n_outputs after loading. Saved model has "
+                        f"n_outputs={saved_n_outputs}, but "
+                        f"n_outputs={requested_n_outputs} was requested."
+                    ) from None
+
+            return model

{braindecode-1.3.2.dev182544824 → braindecode-1.3.2.dev183191339}/braindecode/models/bendr.py

@@ -291,15 +291,23 @@ class BENDR(EEGModuleMixin, nn.Module):
         else:
             in_features = encoder_h
 
+        self._head_in_features = in_features
         self.final_layer = None
         if self.include_final_layer:
-            # in_features: encoder_h*4 (encoder_only) or encoder_h (full model)
-            linear = nn.Linear(in_features=in_features, out_features=self.n_outputs)
-            self.final_layer = nn.utils.parametrizations.weight_norm(
-                linear, name="weight", dim=1
-            )
+            self._build_head(self.n_outputs)
 
-    def forward(self, x):
+    def _build_head(self, n_outputs):
+        linear = nn.Linear(in_features=self._head_in_features, out_features=n_outputs)
+        self.final_layer = nn.utils.parametrizations.weight_norm(
+            linear, name="weight", dim=1
+        )
+
+    def reset_head(self, n_outputs):
+        self._n_outputs = n_outputs
+        self.include_final_layer = True
+        self._build_head(n_outputs)
+
+    def forward(self, x, return_features=False):
         if self.channel_projection is not None:
             x = self.channel_projection(x)
         encoded = self.encoder(x)
@@ -328,6 +336,9 @@ class BENDR(EEGModuleMixin, nn.Module):
             feature = context[:, :, 0]
             # feature: [batch_size, encoder_h]
 
+        if return_features:
+            return {"features": feature, "cls_token": None}
+
         if self.final_layer is not None:
             feature = self.final_layer(feature)
             # feature: [batch_size, n_outputs]

{braindecode-1.3.2.dev182544824 → braindecode-1.3.2.dev183191339}/braindecode/models/biot.py

@@ -183,13 +183,22 @@ class BIOT(EEGModuleMixin, nn.Module):
             attn_layer_dropout=att_layer_drop_prob,
         )
 
+        self._head_activation = activation
         self.final_layer = _ClassificationHead(
             emb_size=self.embed_dim,
             n_outputs=self.n_outputs,
             activation=activation,
         )
 
-    def forward(self, x):
+    def reset_head(self, n_outputs):
+        self._n_outputs = n_outputs
+        self.final_layer = _ClassificationHead(
+            emb_size=self.embed_dim,
+            n_outputs=n_outputs,
+            activation=self._head_activation,
+        )
+
+    def forward(self, x, return_features=False):
         """
         Pass the input through the BIOT encoder, and then through the
         classification head.
@@ -198,15 +207,24 @@ class BIOT(EEGModuleMixin, nn.Module):
         ----------
         x: Tensor
             (batch_size, n_channels, n_times)
+        return_features : bool
+            If True, return a dict with ``"features"`` and ``"cls_token"``
+            instead of the classification output.
 
         Returns
         -------
-        out: Tensor
-            (batch_size, n_outputs)
-        (out, emb): tuple Tensor
-            (batch_size, n_outputs), (batch_size, emb_size)
+        torch.Tensor or tuple or dict
+            Default: ``torch.Tensor`` of shape ``(batch_size, n_outputs)``.
+            If ``return_features=True``: ``dict`` with ``"features"``
+            ``(batch_size, emb_size)`` and ``"cls_token"`` (``None``).
+            If legacy ``return_feature=True`` (init param):
+            ``(out, emb)`` tuple (ignored when ``return_features=True``).
         """
         emb = self.encoder(x)
+
+        if return_features:
+            return {"features": emb, "cls_token": None}
+
         x = self.final_layer(emb)
 
         if self.return_feature:

{braindecode-1.3.2.dev182544824 → braindecode-1.3.2.dev183191339}/braindecode/models/cbramod.py

@@ -205,20 +205,15 @@ class CBraMod(EEGModuleMixin, nn.Module):
         self.encoder = _TransformerEncoder(encoder_layer, num_layers=n_layer)
         self.proj_out = nn.Sequential(nn.Linear(d_model, emb_dim))
 
+        self._emb_dim = emb_dim
+        self._patch_size = patch_size
         self._weights_init()
 
-        try:
-            n_times = self.n_times
-            n_chans = self.n_chans
-        except ValueError:
-            n_times = None
-            n_chans = None
-
         if return_encoder_output:
             self.final_layer = nn.Identity()
-        elif n_times is not None and n_chans is not None:
-            n_patch = n_times // patch_size
-            flat_dim = n_chans * n_patch * emb_dim
+        elif self._n_times is not None and self._n_chans is not None:
+            n_patch = self._n_times // patch_size
+            flat_dim = self._n_chans * n_patch * emb_dim
             self.final_layer = nn.Sequential(
                 nn.Flatten(), nn.Linear(flat_dim, self.n_outputs)
             )
@@ -227,6 +222,17 @@ class CBraMod(EEGModuleMixin, nn.Module):
                 nn.Flatten(), nn.LazyLinear(self.n_outputs)
             )
 
+    def reset_head(self, n_outputs):
+        self._n_outputs = n_outputs
+        if self._n_times is not None and self._n_chans is not None:
+            n_patch = self._n_times // self._patch_size
+            flat_dim = self._n_chans * n_patch * self._emb_dim
+            self.final_layer = nn.Sequential(
+                nn.Flatten(), nn.Linear(flat_dim, n_outputs)
+            )
+        else:
+            self.final_layer = nn.Sequential(nn.Flatten(), nn.LazyLinear(n_outputs))
+
     def _weights_init(self):
         for m in self.modules():
             if isinstance(m, nn.Linear):
@@ -237,11 +243,13 @@ class CBraMod(EEGModuleMixin, nn.Module):
                 nn.init.constant_(m.weight, 1)
                 nn.init.constant_(m.bias, 0)
 
-    def forward(self, x, mask=None):
+    def forward(self, x, mask=None, return_features=False):
         x = self.rearrange(x)
         patch_emb = self.patch_embedding(x, mask)
         feats = self.encoder(patch_emb)
         out = self.proj_out(feats)
+        if return_features:
+            return {"features": out, "cls_token": None}
         return self.final_layer(out)
 
 

{braindecode-1.3.2.dev182544824 → braindecode-1.3.2.dev183191339}/braindecode/models/eegpt.py

@@ -368,6 +368,13 @@ class EEGPT(EEGModuleMixin, nn.Module):
         """Number of temporal patches from encoder."""
         return self.target_encoder.num_patches[1]
 
+    def reset_head(self, n_outputs):
+        self._n_outputs = n_outputs
+        self.final_layer = _LinearConstraintProbe(
+            **self.get_probe_params(),
+            n_outputs=n_outputs,
+        )
+
     def get_probe_params(self) -> dict:
         """Get parameters needed to create a _LinearConstraintProbe.
 
@@ -382,7 +389,7 @@ class EEGPT(EEGModuleMixin, nn.Module):
             "embed_dim": self.embed_dim,
         }
 
-    def forward(self, x):
+    def forward(self, x, return_features=False):
         """
         Forward pass.
 
@@ -390,10 +397,13 @@ class EEGPT(EEGModuleMixin, nn.Module):
         ----------
         x : torch.Tensor
             EEG data of shape (batch, n_chans, n_times).
+        return_features : bool
+            If True, return a dict with ``"features"`` and ``"cls_token"``
+            instead of the classification output.
 
         Returns
         -------
-        torch.Tensor
+        torch.Tensor or dict
             Model output. Shape depends on `n_outputs` and `return_encoder_output`.
         """
         # Channel projection (if configured)
@@ -402,6 +412,9 @@ class EEGPT(EEGModuleMixin, nn.Module):
         # z shape: (batch, n_patches, embed_num, embed_dim)
         z = self.target_encoder(x, self.chans_id)
 
+        if return_features:
+            return {"features": z.flatten(2), "cls_token": None}
+
         if self.return_encoder_output:
             return z
 

{braindecode-1.3.2.dev182544824 → braindecode-1.3.2.dev183191339}/braindecode/models/labram.py

@@ -790,6 +790,7 @@ class Labram(EEGModuleMixin, nn.Module):
         ch_names: list[str] | None = None,
         return_patch_tokens=False,
         return_all_tokens=False,
+        return_features=False,
     ):
         """
         Forward the input EEG data through the model.
@@ -805,6 +806,9 @@ class Labram(EEGModuleMixin, nn.Module):
             If not provided, the channels provided during model initialization
             (via ``chs_info``), channels will be used instead
             If neither is provided, an error will be raised.
+        return_features : bool
+            If True, return a dict with ``"features"`` (patch tokens) and
+            ``"cls_token"`` instead of the classification output.
         return_patch_tokens: bool
             Return the patch tokens
         return_all_tokens: bool
@@ -812,12 +816,18 @@ class Labram(EEGModuleMixin, nn.Module):
 
         Returns
         -------
-        torch.Tensor
+        torch.Tensor or dict
             The output of the model with dimensions (batch, n_outputs)
         """
         # Apply automatic channel reordering if configured and input_chans not provided
         x, input_chans = self._select_channels(x, ch_names=ch_names)
 
+        if return_features:
+            x = self.forward_features(
+                x, input_chans=input_chans, return_all_tokens=True
+            )
+            return {"features": x[:, 1:, :], "cls_token": x[:, 0, :]}
+
         x = self.forward_features(
             x,
             input_chans=input_chans,
@@ -851,6 +861,10 @@ class Labram(EEGModuleMixin, nn.Module):
             nn.Linear(self.embed_dim, n_outputs) if n_outputs > 0 else nn.Identity()
         )
 
+    def reset_head(self, n_outputs):
+        self._n_outputs = n_outputs
+        self.reset_classifier(n_outputs)
+
     def _adj_temporal_embedding(self, num_ch, batch_size, dim_embed=None):
         """
         Adjust the dimensions of the time embedding to match the

{braindecode-1.3.2.dev182544824 → braindecode-1.3.2.dev183191339}/braindecode/models/luna.py

@@ -49,7 +49,7 @@ class LUNA(EEGModuleMixin, nn.Module):
        **Pre-trained Weights Available**
 
        This model has pre-trained weights available on the Hugging Face Hub
-       at `thorir/LUNA <https://huggingface.co/thorir/LUNA>`_.
+       at `PulpBio/LUNA <https://huggingface.co/PulpBio/LUNA>`_.
 
        Available model variants:
 
@@ -61,19 +61,13 @@ class LUNA(EEGModuleMixin, nn.Module):
 
        .. code-block:: python
 
-           from huggingface_hub import hf_hub_download
-           from safetensors.torch import load_file
            from braindecode.models import LUNA
 
-           # Download pre-trained weights
-           model_path = hf_hub_download(
-               repo_id="thorir/LUNA",
+           # Load pre-trained base model from Hugging Face Hub
+           model = LUNA.from_pretrained(
+               "PulpBio/LUNA",
                filename="LUNA_base.safetensors",
-           )
-
-           # Create model for classification (fine-tuning)
-           model = LUNA(
-               n_outputs=2,  # Number of classes for your task
+               n_outputs=2,
                n_chans=22,
                n_times=1000,
                embed_dim=64,
@@ -81,11 +75,6 @@ class LUNA(EEGModuleMixin, nn.Module):
                depth=8,
            )
 
-           # Load pre-trained encoder weights
-           state_dict = load_file(model_path)
-           # load_state_dict applies model.mapping automatically
-           model.load_state_dict(state_dict, strict=False)
-
        To push your own trained model to the Hub:
 
        .. code-block:: python

{braindecode-1.3.2.dev182544824 → braindecode-1.3.2.dev183191339}/braindecode/models/reve.py

@@ -311,19 +311,7 @@ class REVE(EEGModuleMixin, nn.Module):
             geglu=self.use_geglu,
         )
 
-        final_dim = self._get_flattened_output_dim()
-
-        if self.use_attention_pooling:
-            self.final_layer = nn.Sequential(
-                nn.LayerNorm(self.embed_dim),
-                nn.Linear(self.embed_dim, self.n_outputs),
-            )
-        else:
-            self.final_layer = nn.Sequential(
-                nn.Flatten(),
-                nn.LayerNorm(final_dim),
-                nn.Linear(final_dim, self.n_outputs),
-            )
+        self._build_head(self.n_outputs)
 
         if self.use_attention_pooling:
             self.cls_query_token = nn.Parameter(torch.randn(1, 1, self.embed_dim))
@@ -352,6 +340,24 @@ class REVE(EEGModuleMixin, nn.Module):
         flat_dim = self.n_chans * n_patches * self.embed_dim
         return flat_dim
 
+    def _build_head(self, n_outputs):
+        if self.use_attention_pooling:
+            self.final_layer = nn.Sequential(
+                nn.LayerNorm(self.embed_dim),
+                nn.Linear(self.embed_dim, n_outputs),
+            )
+        else:
+            final_dim = self._get_flattened_output_dim()
+            self.final_layer = nn.Sequential(
+                nn.Flatten(),
+                nn.LayerNorm(final_dim),
+                nn.Linear(final_dim, n_outputs),
+            )
+
+    def reset_head(self, n_outputs):
+        self._n_outputs = n_outputs
+        self._build_head(n_outputs)
+
     def get_positions(self, channel_names: list[str]) -> torch.Tensor:
         """Fetch channel positions from the position bank. The position bank is downloaded when the model is instantiated.
 
@@ -373,7 +379,8 @@ class REVE(EEGModuleMixin, nn.Module):
         eeg: torch.Tensor,
         pos: Optional[torch.Tensor] = None,
         return_output: bool = False,
-    ) -> Union[torch.Tensor, list[torch.Tensor]]:
+        return_features: bool = False,
+    ) -> Union[torch.Tensor, list[torch.Tensor], dict]:
         """
         Forward pass of the model.
 
@@ -392,17 +399,23 @@ class REVE(EEGModuleMixin, nn.Module):
         return_output : bool, optional
             If True, returns the output from the transformer directly.
             If False, applies the final layer and returns the processed output. Default is False.
+        return_features : bool, optional
+            If True, returns a dict ``{"features": Tensor, "cls_token": None}``
+            with encoder features before the final layer. Default is False.
 
         Returns
         -------
-        Union[torch.Tensor, list[torch.Tensor]]
-            - If `return_output` is False: Returns a single `torch.Tensor` (output after final layer).
-            - If `return_output` is True: Returns a `list[torch.Tensor]` (outputs from transformer layers).
-
-            The output tensor(s) from the model. If `return_output` is True,
-            returns the transformer output; otherwise, returns the output after the final layer.
+        torch.Tensor or list[torch.Tensor] or dict
+            Default: ``torch.Tensor`` after the final layer.
+            If ``return_output=True``: ``list[torch.Tensor]`` from transformer layers.
+            If ``return_features=True``: ``dict`` with ``"features"`` and ``"cls_token"`` keys.
         """
 
+        if return_output and return_features:
+            raise ValueError(
+                "return_output and return_features are mutually exclusive."
+            )
+
         patches = eeg.unfold(
             dimension=2,
             size=self.patch_size,
@@ -453,6 +466,9 @@ class REVE(EEGModuleMixin, nn.Module):
             emb=self.embed_dim,
         )
 
+        if return_features:
+            return {"features": x, "cls_token": None}
+
         if self.use_attention_pooling:
             x = self._attention_pooling(x)