Braindecode 1.3.2.dev183566201__tar.gz → 1.3.2.dev183987088__tar.gz

{braindecode-1.3.2.dev183566201/braindecode.egg-info → braindecode-1.3.2.dev183987088}/PKG-INFO

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

@@ -150,40 +150,14 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
         # Load pretrained model
         model = {name}.from_pretrained("username/my-{name_lower}-model")
 
-        # 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)
+    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.
 
     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):
@@ -468,34 +442,6 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
         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):
@@ -695,78 +641,15 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
             **model_kwargs,
         ):
             model_kwargs.pop("braindecode_version", None)
-            filename = model_kwargs.pop("filename", None)
             resolve_type_kwargs(cls, 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
+            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,
+            )

{braindecode-1.3.2.dev183566201 → braindecode-1.3.2.dev183987088}/braindecode/models/bendr.py

@@ -291,23 +291,15 @@ 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:
-            self._build_head(self.n_outputs)
-
-    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)
+            # 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
+            )
 
-    def forward(self, x, return_features=False):
+    def forward(self, x):
         if self.channel_projection is not None:
             x = self.channel_projection(x)
         encoded = self.encoder(x)
@@ -336,9 +328,6 @@ 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.dev183566201 → braindecode-1.3.2.dev183987088}/braindecode/models/biot.py

@@ -183,22 +183,13 @@ 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 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):
+    def forward(self, x):
         """
         Pass the input through the BIOT encoder, and then through the
         classification head.
@@ -207,24 +198,15 @@ 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
         -------
-        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``).
+        out: Tensor
+            (batch_size, n_outputs)
+        (out, emb): tuple Tensor
+            (batch_size, n_outputs), (batch_size, emb_size)
         """
         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.dev183566201 → braindecode-1.3.2.dev183987088}/braindecode/models/cbramod.py

@@ -205,15 +205,20 @@ 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 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
+        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
             self.final_layer = nn.Sequential(
                 nn.Flatten(), nn.Linear(flat_dim, self.n_outputs)
             )
@@ -222,17 +227,6 @@ 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):
@@ -243,13 +237,11 @@ class CBraMod(EEGModuleMixin, nn.Module):
                 nn.init.constant_(m.weight, 1)
                 nn.init.constant_(m.bias, 0)
 
-    def forward(self, x, mask=None, return_features=False):
+    def forward(self, x, mask=None):
         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.dev183566201 → braindecode-1.3.2.dev183987088}/braindecode/models/eeginception_mi.py

@@ -1,5 +1,4 @@
 # Authors: Cedric Rommel <cedric.rommel@inria.fr>
-#          Sarthak Tayal <sarthaktayal2@gmail.com>
 #
 # License: BSD (3-clause)
 
@@ -96,7 +95,7 @@ class EEGInceptionMI(EEGModuleMixin, nn.Module):
 
         self.mapping = {
             "fc.weight": "final_layer.fc.weight",
-            "fc.bias": "final_layer.fc.bias",
+            "tc.bias": "final_layer.fc.bias",
         }
 
         # ======== Inception branches ========================

{braindecode-1.3.2.dev183566201 → braindecode-1.3.2.dev183987088}/braindecode/models/eegitnet.py

@@ -1,5 +1,4 @@
 # Authors: Ghaith Bouallegue <ghaithbouallegue@gmail.com>
-#          Sarthak Tayal <sarthaktayal2@gmail.com>
 #
 # License: BSD-3
 from einops.layers.torch import Rearrange
@@ -95,8 +94,8 @@ class EEGITNet(EEGModuleMixin, nn.Sequential):
             sfreq=sfreq,
         )
         self.mapping = {
-            "classification.1.weight": "final_layer.weight",
-            "classification.1.bias": "final_layer.bias",
+            "classification.1.weight": "final_layer.clf.weight",
+            "classification.1.bias": "final_layer.clf.weight",
         }
 
         del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
@@ -122,7 +121,7 @@ class EEGITNet(EEGModuleMixin, nn.Sequential):
         block13 = self._get_inception_branch(
             in_channels=self.n_chans,
             out_channels=n_filters_time * 4,
-            kernel_length=kernel_length * 4,
+            kernel_length=n_filters_time * 4,
             activation=activation,
         )
         self.add_module("inception_block", InceptionBlock((block11, block12, block13)))

{braindecode-1.3.2.dev183566201 → braindecode-1.3.2.dev183987088}/braindecode/models/eegpt.py

@@ -368,13 +368,6 @@ 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.
 
@@ -389,7 +382,7 @@ class EEGPT(EEGModuleMixin, nn.Module):
             "embed_dim": self.embed_dim,
         }
 
-    def forward(self, x, return_features=False):
+    def forward(self, x):
         """
         Forward pass.
 
@@ -397,13 +390,10 @@ 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 or dict
+        torch.Tensor
             Model output. Shape depends on `n_outputs` and `return_encoder_output`.
         """
         # Channel projection (if configured)
@@ -412,9 +402,6 @@ 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.dev183566201 → braindecode-1.3.2.dev183987088}/braindecode/models/labram.py

@@ -790,7 +790,6 @@ 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.
@@ -806,9 +805,6 @@ 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
@@ -816,18 +812,12 @@ class Labram(EEGModuleMixin, nn.Module):
 
         Returns
         -------
-        torch.Tensor or dict
+        torch.Tensor
             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,
@@ -861,10 +851,6 @@ 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.dev183566201 → braindecode-1.3.2.dev183987088}/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 `PulpBio/LUNA <https://huggingface.co/PulpBio/LUNA>`_.
+       at `thorir/LUNA <https://huggingface.co/thorir/LUNA>`_.
 
        Available model variants:
 
@@ -61,13 +61,19 @@ 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
 
-           # Load pre-trained base model from Hugging Face Hub
-           model = LUNA.from_pretrained(
-               "PulpBio/LUNA",
+           # Download pre-trained weights
+           model_path = hf_hub_download(
+               repo_id="thorir/LUNA",
                filename="LUNA_base.safetensors",
-               n_outputs=2,
+           )
+
+           # Create model for classification (fine-tuning)
+           model = LUNA(
+               n_outputs=2,  # Number of classes for your task
                n_chans=22,
                n_times=1000,
                embed_dim=64,
@@ -75,6 +81,11 @@ 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.dev183566201 → braindecode-1.3.2.dev183987088}/braindecode/models/reve.py

@@ -311,7 +311,19 @@ class REVE(EEGModuleMixin, nn.Module):
             geglu=self.use_geglu,
         )
 
-        self._build_head(self.n_outputs)
+        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),
+            )
 
         if self.use_attention_pooling:
             self.cls_query_token = nn.Parameter(torch.randn(1, 1, self.embed_dim))
@@ -340,24 +352,6 @@ 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.
 
@@ -379,8 +373,7 @@ class REVE(EEGModuleMixin, nn.Module):
         eeg: torch.Tensor,
         pos: Optional[torch.Tensor] = None,
         return_output: bool = False,
-        return_features: bool = False,
-    ) -> Union[torch.Tensor, list[torch.Tensor], dict]:
+    ) -> Union[torch.Tensor, list[torch.Tensor]]:
         """
         Forward pass of the model.
 
@@ -399,22 +392,16 @@ 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
         -------
-        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.
-        """
+        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).
 
-        if return_output and return_features:
-            raise ValueError(
-                "return_output and return_features are mutually exclusive."
-            )
+            The output tensor(s) from the model. If `return_output` is True,
+            returns the transformer output; otherwise, returns the output after the final layer.
+        """
 
         patches = eeg.unfold(
             dimension=2,
@@ -466,9 +453,6 @@ 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)