braindecode 1.8.1.dev169692045__tar.gz → 1.8.1.dev169698056__tar.gz

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169698056}/NOTICE.txt

@@ -47,7 +47,6 @@ The following components are licensed under the MIT License:
 - `braindecode/models/mirepnet.py`
 - `braindecode/models/popt.py`
 - `braindecode/models/tfm_tokenizer.py`
-- `braindecode/modules/popt_modules.py`
 - `braindecode/models/brainomni.py`
 - `braindecode/modules/quantization.py`
 - `braindecode/modules/linear_attention.py`

{braindecode-1.8.1.dev169692045/braindecode.egg-info → braindecode-1.8.1.dev169698056}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.8.1.dev169692045
+Version: 1.8.1.dev169698056
 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>
@@ -32,8 +32,8 @@ Requires-Dist: joblib
 Requires-Dist: torchinfo
 Requires-Dist: pandas
 Requires-Dist: wfdb>=4.3.1
-Requires-Dist: linear_attention_transformer
 Requires-Dist: docstring_inheritance
+Requires-Dist: einops
 Requires-Dist: pydantic>=2.0
 Provides-Extra: moabb
 Requires-Dist: moabb>=1.4.3; extra == "moabb"

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169698056}/braindecode/augmentation/functional.py

@@ -1089,10 +1089,12 @@ def mixup(
         Online: https://arxiv.org/abs/1710.09412
     .. [2] https://github.com/facebookresearch/mixup-cifar10/blob/master/train.py
     """
-    device = X.device
     batch_size, n_channels, n_times = X.shape
 
-    X_mix = torch.zeros((batch_size, n_channels, n_times)).to(device)
+    X_mix = X.new_zeros(
+        (batch_size, n_channels, n_times),
+        dtype=X.dtype if X.is_floating_point() else torch.get_default_dtype(),
+    )
     y_a = y.clone()
     y_b = y[idx_perm].clone()
 

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169698056}/braindecode/datasets/chb_mit.py

@@ -16,12 +16,7 @@ It reorganizes the file structure to comply with the BIDS specification. To this
 # License: BSD (3-clause)
 from __future__ import annotations
 
-from pathlib import Path
-
-from mne.datasets import fetch_dataset
-
-from braindecode.datasets import BIDSDataset
-from braindecode.datasets.utils import _correct_dataset_path
+from braindecode.datasets.utils import _ZenodoBIDSDataset
 
 CHB_MIT_URL = "https://zenodo.org/records/10259996/files/BIDS_CHB-MIT.zip"
 CHB_MIT_archive_name = "chb_mit_bids.zip"
@@ -38,7 +33,7 @@ CHB_MIT_dataset_params = {
 }
 
 
-class CHBMIT(BIDSDataset):
+class CHBMIT(_ZenodoBIDSDataset):
     """The Children's Hospital Boston EEG Dataset.
 
     This database, collected at the Children's Hospital Boston, consists of EEG recordings
@@ -129,35 +124,6 @@ class CHBMIT(BIDSDataset):
     .. footbibliography::
     """
 
-    def __init__(self, root=None, *args, **kwargs):
-        # Download dataset if not present
-        if root is None:
-            path_root = fetch_dataset(
-                dataset_params=CHB_MIT_dataset_params,
-                path=None,
-                processor="unzip",
-                force_update=False,
-            )
-            # First time we fetch the dataset, we need to move the files to the
-            # correct directory.
-            path_root = _correct_dataset_path(
-                path_root, CHB_MIT_archive_name, "BIDS_CHB-MIT"
-            )
-        else:
-            # Validate that the provided root is a valid BIDS dataset
-            if not Path(f"{root}/participants.tsv").exists():
-                raise ValueError(
-                    f"The provided root directory {root} does not contain a valid "
-                    "BIDS dataset (missing participants.tsv). Please ensure the "
-                    "root points directly to the BIDS dataset directory."
-                )
-            path_root = root
-
-        kwargs["root"] = path_root
-
-        super().__init__(
-            *args,
-            extensions=".edf",
-            check=False,
-            **kwargs,
-        )
+    _dataset_params = CHB_MIT_dataset_params
+    _archive_name = CHB_MIT_archive_name
+    _unzipped_dir = "BIDS_CHB-MIT"

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169698056}/braindecode/datasets/siena.py

@@ -15,12 +15,7 @@ It reorganizes the file structure to comply with the BIDS specification. To this
 # License: BSD (3-clause)
 from __future__ import annotations
 
-from pathlib import Path
-
-from mne.datasets import fetch_dataset
-
-from braindecode.datasets import BIDSDataset
-from braindecode.datasets.utils import _correct_dataset_path
+from braindecode.datasets.utils import _ZenodoBIDSDataset
 
 SIENA_URL = "https://zenodo.org/records/10640762/files/BIDS_Siena.zip"
 SIENA_archive_name = "SIENA.zip"
@@ -37,7 +32,7 @@ SIENA_dataset_params = {
 }
 
 
-class SIENA(BIDSDataset):
+class SIENA(_ZenodoBIDSDataset):
     """The Siena EEG Dataset.
 
     The database consists of EEG recordings of 14 patients acquired at the Unit of Neurology
@@ -128,35 +123,6 @@ class SIENA(BIDSDataset):
         automated seizure detection algorithms. Epilepsia, 66, 14-24.
     """
 
-    def __init__(self, root=None, *args, **kwargs):
-        # Download dataset if not present
-        if root is None:
-            path_root = fetch_dataset(
-                dataset_params=SIENA_dataset_params,
-                path=None,
-                processor="unzip",
-                force_update=False,
-            )
-            # First time we fetch the dataset, we need to move the files to the
-            # correct directory.
-            path_root = _correct_dataset_path(
-                path_root, SIENA_archive_name, "BIDS_Siena"
-            )
-        else:
-            # Validate that the provided root is a valid BIDS dataset
-            if not Path(f"{root}/participants.tsv").exists():
-                raise ValueError(
-                    f"The provided root directory {root} does not contain a valid "
-                    "BIDS dataset (missing participants.tsv). Please ensure the "
-                    "root points directly to the BIDS dataset directory."
-                )
-            path_root = root
-
-        kwargs["root"] = path_root
-
-        super().__init__(
-            *args,
-            extensions=".edf",
-            check=False,
-            **kwargs,
-        )
+    _dataset_params = SIENA_dataset_params
+    _archive_name = SIENA_archive_name
+    _unzipped_dir = "BIDS_Siena"

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169698056}/braindecode/datasets/sleep_physio_challe_18.py

@@ -57,52 +57,6 @@ def ensure_metafiles_exist():
         _download_if_missing(file_path, url)
 
 
-def _update_pc18_sleep_records(fname=PC18_RECORDS):
-    """Create CSV file with information about available PC18 recordings."""
-    # Load and massage the checksums.
-    sha_train_df = pd.read_csv(
-        PC18_SHA1_TRAINING,
-        sep="  ",
-        header=None,
-        names=["sha", "fname"],
-        engine="python",
-    )
-    sha_test_df = pd.read_csv(
-        PC18_SHA1_TEST, sep="  ", header=None, names=["sha", "fname"], engine="python"
-    )
-    sha_train_df["Split"] = "training"
-    sha_test_df["Split"] = "test"
-    sha_df = pd.concat([sha_train_df, sha_test_df], axis=0, ignore_index=True)
-    select_records = (
-        sha_df.fname.str.startswith("tr") | sha_df.fname.str.startswith("te")
-    ) & ~sha_df.fname.str.endswith("arousal.mat")
-    sha_df = sha_df[select_records]
-    sha_df["Record"] = sha_df["fname"].str.split("/", expand=True)[0]
-    sha_df["fname"] = sha_df[["Split", "fname"]].agg("/".join, axis=1)
-
-    # Load and massage the data.
-    data = pd.read_csv(PC18_INFO)
-
-    data = data.reset_index().rename({"index": "Subject"}, axis=1)
-    data["Sex"] = (
-        data["Sex"].map({"F": "female", "M": "male", "m": "male"}).astype("category")
-    )
-    data = sha_df.merge(data, on="Record")
-
-    data["Record type"] = (
-        data["fname"]
-        .str.split(".", expand=True)[1]
-        .map({"hea": "Header", "mat": "PSG", "arousal": "Arousal"})
-        .astype("category")
-    )
-    data = data[
-        ["Subject", "Record", "Record type", "Split", "Age", "Sex", "sha", "fname"]
-    ].sort_values(by="Subject")
-
-    # Save the data.
-    data.to_csv(fname, index=False)
-
-
 def _data_path(path=None):
     """Get path to local copy of PC18 dataset."""
     key = "PC18_DATASET_PATH"

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169698056}/braindecode/datasets/utils.py

@@ -9,6 +9,10 @@ from __future__ import annotations
 import os
 from pathlib import Path
 
+from mne.datasets import fetch_dataset
+
+from .bids import BIDSDataset
+
 
 def _correct_dataset_path(
     path: str, archive_name: str, subfolder_name: str | None = None
@@ -65,3 +69,36 @@ def _correct_dataset_path(
             path = subfolder_path
 
     return path
+
+
+class _ZenodoBIDSDataset(BIDSDataset):
+    """A BIDS dataset in EDF files, downloaded from a Zenodo archive when no
+    ``root`` is given. Subclasses set the three class attributes below."""
+
+    _dataset_params: dict
+    _archive_name: str  # the downloaded archive
+    _unzipped_dir: str  # the folder the archive extracts to
+
+    def __init__(self, root=None, *args, **kwargs):
+        if root is None:
+            path_root = fetch_dataset(
+                dataset_params=self._dataset_params,
+                path=None,
+                processor="unzip",
+                force_update=False,
+            )
+            # The first fetch extracts into a nested folder: point at the data.
+            path_root = _correct_dataset_path(
+                path_root, self._archive_name, self._unzipped_dir
+            )
+        elif not Path(f"{root}/participants.tsv").exists():
+            raise ValueError(
+                f"The provided root directory {root} does not contain a valid "
+                "BIDS dataset (missing participants.tsv). Please ensure the "
+                "root points directly to the BIDS dataset directory."
+            )
+        else:
+            path_root = root
+        super().__init__(
+            *args, root=path_root, extensions=".edf", check=False, **kwargs
+        )

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169698056}/braindecode/functional/initialization.py

@@ -1,25 +1,30 @@
 import math
 
+import torch
 from torch import nn
 
 
 def glorot_weight_zero_bias(model):
-    """Initialize parameters of all modules by initializing weights with.
+    """Set the weights of BatchNorm layers to one and all biases to zero.
 
-    glorot uniform/xavier initialization, and setting biases to zero. Weights from
-    batch norm layers are set to 1.
+    Despite its name, this function does not apply the Glorot (Xavier)
+    initialization: other weights keep the initialization of their layer
+    (PyTorch's default, or Xavier uniform for the max-norm constrained layers
+    such as :class:`braindecode.modules.Conv2dWithConstraint`). The name is
+    kept for backward compatibility.
 
     Parameters
     ----------
     model : Module
     """
     for module in model.modules():
-        if hasattr(module, "weight"):
-            if "BatchNorm" in module.__class__.__name__:
-                nn.init.constant_(module.weight, 1)
-        if hasattr(module, "bias"):
-            if module.bias is not None:
-                nn.init.constant_(module.bias, 0)
+        # ``weight``/``bias`` can be None (e.g. BatchNorm with affine=False)
+        weight = getattr(module, "weight", None)
+        if isinstance(weight, torch.Tensor) and "BatchNorm" in type(module).__name__:
+            nn.init.constant_(weight, 1)
+        bias = getattr(module, "bias", None)
+        if isinstance(bias, torch.Tensor):
+            nn.init.constant_(bias, 0)
 
 
 def rescale_parameter(param, layer_id):

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169698056}/braindecode/models/brainbert.py

@@ -8,6 +8,8 @@ the STFT front-end moved inside the model. See :class:`BrainBERT`.
 
 from __future__ import annotations
 
+import warnings
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -16,6 +18,11 @@ from einops.layers.torch import Rearrange, Reduce
 from braindecode.functional import sinusoidal_positional_encoding, spectral_input
 from braindecode.models.base import EEGModuleMixin
 from braindecode.models.util import warn_if_sfreq_differs
+from braindecode.modules.blocks import (
+    _check_on_non_divisible,
+    _pad_or_crop_to_patches,
+    _patch_count,
+)
 
 
 class BrainBERT(EEGModuleMixin, nn.Module, license="unknown"):
@@ -122,6 +129,14 @@ class BrainBERT(EEGModuleMixin, nn.Module, license="unknown"):
         pretrained).
     drop_prob : float, optional
         Dropout probability. Default 0.1.
+    on_non_divisible : {"pad", "crop", "error"}, optional
+        A window too short to give the ``pool_n_frames`` frames the pooling
+        needs (``n_times < 1401`` with the defaults). ``"pad"`` (default)
+        adds the missing STFT frames as zeros, half before and half after the
+        window, with a warning at construction: the ``stft_clip`` frames at
+        each end are trimmed and the centre ones pooled, so the signal stays
+        in the pooled frames. ``"crop"`` and ``"error"`` raise. A long enough
+        window gives the same output under all three.
 
     References
     ----------
@@ -147,6 +162,7 @@ class BrainBERT(EEGModuleMixin, nn.Module, license="unknown"):
         pool_n_frames: int | None = 10,
         activation: type[nn.Module] = nn.GELU,
         drop_prob: float = 0.1,
+        on_non_divisible: str = "pad",
         # --- braindecode mandatory signal parameters ---
         n_outputs=None,
         n_chans=None,
@@ -176,6 +192,7 @@ class BrainBERT(EEGModuleMixin, nn.Module, license="unknown"):
                 f"pool_n_frames must be positive; got {pool_n_frames}. "
                 "Pass None to average every frame instead."
             )
+        _check_on_non_divisible(on_non_divisible)
         self.pool_n_frames = pool_n_frames
         self.min_frames = 1 if pool_n_frames is None else pool_n_frames
         warn_if_sfreq_differs("BrainBERT", self._sfreq, 2048)
@@ -186,15 +203,22 @@ class BrainBERT(EEGModuleMixin, nn.Module, license="unknown"):
             idx_freq_cutoff=idx_freq_cutoff,
             clip=stft_clip,
             zscore_before_clip=stft_zscore_before_clip,
+            min_frames=self.min_frames if on_non_divisible == "pad" else 0,
         )
         n_frames = self.spectrogram.n_frames(self.n_times)
         if n_frames < self.min_frames:
             min_n_times = self.spectrogram.min_n_times(self.min_frames)
-            raise ValueError(
+            message = (
                 f"n_times={self.n_times} gives {n_frames} STFT frames; the pooling "
-                f"needs {self.min_frames}, i.e. n_times >= {min_n_times} with these "
-                "STFT settings."
+                f"needs {self.min_frames}, i.e. n_times >= {min_n_times} with "
+                "these STFT settings"
             )
+            if on_non_divisible != "pad":
+                raise ValueError(message + ".")
+            warnings.warn(
+                message + ": the input will be padded.", UserWarning, stacklevel=2
+            )
+            n_frames = self.min_frames
 
         # The re-hosted Hub checkpoint already uses this port's names. The
         # authors' original ``.pth`` calls the input block ``input_encoding``;
@@ -329,6 +353,10 @@ class _STFTSpectrogram(nn.Module):
         Number of boundary frames trimmed from each end.
     zscore_before_clip : bool
         Whether the per-bin z-score over time precedes the trimming.
+    min_frames : int
+        A signal giving fewer output frames gets the missing STFT frames as
+        zeros, half before and half after it. ``0`` (default): scipy's framing
+        only.
     """
 
     def __init__(
@@ -338,6 +366,7 @@ class _STFTSpectrogram(nn.Module):
         idx_freq_cutoff: int = 40,
         clip: int = 10,
         zscore_before_clip: bool = True,
+        min_frames: int = 0,
     ):
         super().__init__()
         if not 0 <= noverlap < nperseg:
@@ -355,6 +384,7 @@ class _STFTSpectrogram(nn.Module):
         self.idx_freq_cutoff = idx_freq_cutoff
         self.clip = clip
         self.zscore_before_clip = zscore_before_clip
+        self.min_frames = min_frames
         # "zscore" as upstream; pipelines that apply their own statistics (the
         # Neuroprobe global z-score recipes) set it to "none" on the instance.
         self.normalizing = "zscore"
@@ -365,15 +395,11 @@ class _STFTSpectrogram(nn.Module):
             "window", torch.hann_window(nperseg, periodic=True), persistent=False
         )
 
-    def _padded_length(self, n_times: int) -> int:
-        """Signal length after scipy's ``boundary="zeros"`` and ``padded=True``."""
-        length = n_times + 2 * self.boundary_pad
-        return length + (-(length - self.nperseg)) % self.step
-
     def n_frames(self, n_times: int) -> int:
-        """Number of output frames for a signal of ``n_times`` samples."""
-        n_seg = (self._padded_length(n_times) - self.nperseg) // self.step + 1
-        return n_seg - 2 * self.clip
+        """Number of frames scipy's framing gives ``n_times`` samples: segments of
+        ``nperseg`` every ``step`` (the last one zero-padded), minus the clip."""
+        length = n_times + 2 * self.boundary_pad
+        return _patch_count(length, self.nperseg, "pad", self.step) - 2 * self.clip
 
     def min_n_times(self, n_frames: int) -> int:
         """Shortest signal that yields at least ``n_frames`` output frames."""
@@ -382,15 +408,20 @@ class _STFTSpectrogram(nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """``(batch, n_chans, n_times)`` -> ``(batch, n_chans, n_frames, cutoff)``."""
-        # scipy boundary="zeros": half a window of zeros on each end, then
-        # padded=True: extend so an integer number of segments fits.
         if not x.is_floating_point():
             raise TypeError(
                 f"BrainBERT expects a floating-point signal, got {x.dtype}."
             )
-        n_times = x.shape[-1]
-        right_pad = self._padded_length(n_times) - n_times - self.boundary_pad
-        xp = F.pad(x, (self.boundary_pad, right_pad))  # (batch, n_chans, padded)
+        # scipy boundary="zeros" (half a window of zeros on each end), then
+        # padded=True (right zero-pad so the segments tile the signal). A signal
+        # short of min_frames gets the missing segments as zeros, half before
+        # it: the clip frames at each end are trimmed, so a right pad alone
+        # would leave it outside the frames that are kept.
+        xp = F.pad(x, (self.boundary_pad, self.boundary_pad))
+        min_segments = self.min_frames + 2 * self.clip if self.min_frames else 1
+        xp = _pad_or_crop_to_patches(
+            xp, self.nperseg, "pad", self.step, min_segments, center=True
+        )
         frames = xp.unfold(-1, self.nperseg, self.step)  # (b, c, n_frames, nperseg)
         frames = spectral_input(frames)
         win = self.window.to(frames)

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169698056}/braindecode/models/mapa.py

@@ -27,7 +27,7 @@ from torch import nn
 
 from braindecode.functional import rescale_parameter, rotate_pairs, spectral_input
 from braindecode.models.base import EEGModuleMixin
-from braindecode.models.util import dkt_region_slots
+from braindecode.models.util import _in_model_resample, dkt_region_slots
 from braindecode.modules import FeedForwardBlock
 
 # Blocks whose output carries a deep-supervision norm, following V-JEPA 2.1.
@@ -104,9 +104,11 @@ class MAPA(EEGModuleMixin, nn.Module, license="apache-2.0"):
     batch share them. Only ``pooling="flatten"`` is tied to one montage and one
     window length.
 
-    *Sampling rate and window length.* Resample to 2048 Hz. A window yields
-    ``1 + n_times // 64`` frames, truncated to a multiple of 8, so it needs at
-    least 448 samples. With ``normalization="session"`` the input is the
+    *Sampling rate and window length.* The model reads 2048 Hz; another
+    integer ``sfreq`` is resampled to 2048 Hz inside the model with
+    :class:`torchaudio.transforms.Resample`, with a warning. A 2048 Hz window
+    yields ``1 + n_times // 64`` frames, truncated to a multiple of 8, so it
+    needs at least 448 samples. With ``normalization="session"`` the input is the
     spectrogram itself (``sfreq=32``, ``n_times`` in frames): the 20 retained
     bins, robust z-scored per contact and bin over the whole recording.
 
@@ -252,18 +254,23 @@ class MAPA(EEGModuleMixin, nn.Module, license="apache-2.0"):
                 f"{normalization!r}."
             )
         # The bands are FFT bin slices at 2048 Hz on a 32 Hz frame clock, so
-        # another rate would shift the bands and the token rate; session inputs
-        # are already the 32 Hz spectrogram. sfreq is advisory for the
-        # frontend, so it is validated only when it is known, and a wrong
-        # value is a hard error rather than a silent warning.
+        # another raw rate is resampled to 2048 Hz in forward; session inputs
+        # are already the 32 Hz spectrogram. sfreq is validated only when it
+        # is known.
         expected_sfreq = 32 if normalization == "session" else 2048
         sfreq_known = self._sfreq is not None or (
             self._input_window_seconds is not None and self._n_times is not None
         )
+        self.resample: nn.Module = nn.Identity()
+        n_times_model = self.n_times
         if sfreq_known and float(self.sfreq) != expected_sfreq:
-            raise ValueError(
-                f"MAPA with normalization={normalization!r} needs sfreq="
-                f"{expected_sfreq} Hz, got {self.sfreq} Hz."
+            if normalization == "session":
+                raise ValueError(
+                    f"MAPA with normalization={normalization!r} needs sfreq="
+                    f"{expected_sfreq} Hz, got {self.sfreq} Hz."
+                )
+            self.resample, n_times_model = _in_model_resample(
+                "MAPA", self.sfreq, 2048, self.n_times
             )
 
         self.d_model = d_model
@@ -276,7 +283,7 @@ class MAPA(EEGModuleMixin, nn.Module, license="apache-2.0"):
         self.activation = activation
 
         self.n_frames = _frame_count(
-            self.n_times, spectrogram=normalization == "session"
+            n_times_model, spectrogram=normalization == "session"
         )
         self.k_full = sum(self.n_frames // stride for *_, stride in _BANDS)
 
@@ -617,6 +624,7 @@ class MAPA(EEGModuleMixin, nn.Module, license="apache-2.0"):
                 f"(batch, n_chans, n_times), but got {list(x.shape)} with "
                 f"normalization='{self.normalization}'."
             )
+        x = self.resample(x)
         indices = self._resolve_sensor_indices(sensor_indices, x)
         n_frames = _frame_count(x.shape[-1], spectrogram=spectrogram)
         if self.pooling == "flatten" and (

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169698056}/braindecode/models/medformer.py

@@ -525,24 +525,6 @@ class _AttentionLayer(nn.Module):
         return self.out_projection(out), attn
 
 
-class _TriangularCausalMask:
-    def __init__(
-        self, batch_size: int, seq_len: int, device: Optional[torch.device] = None
-    ):
-        # Normalize device to a torch.device for .to(device)
-        if device is None:
-            device = torch.device("cpu")
-        mask_shape = [batch_size, 1, seq_len, seq_len]
-        with torch.no_grad():
-            self._mask = torch.triu(
-                torch.ones(mask_shape, dtype=torch.bool), diagonal=1
-            ).to(device)
-
-    @property
-    def mask(self) -> torch.Tensor:
-        return self._mask
-
-
 class _FullAttention(nn.Module):
     def __init__(
         self,

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169698056}/braindecode/models/neurorvq.py

@@ -14,11 +14,11 @@ from typing import Dict, Optional, Union
 import torch
 import torch.nn.functional as F
 from torch import Tensor, nn
-from torchaudio.transforms import Resample
 
 from braindecode.models.base import EEGModuleMixin
 from braindecode.models.labram import _Attention
 from braindecode.models.signal_jepa import _GroupNorm
+from braindecode.models.util import _in_model_resample
 from braindecode.modules import MLP, DropPath, PatchTokenizer
 from braindecode.modules.blocks import _check_on_non_divisible, _patch_count
 
@@ -534,23 +534,9 @@ class NeuroRVQ(EEGModuleMixin, nn.Module, license="cc-by-nc-4.0"):
         self.resample: nn.Module = nn.Identity()
         self._n_times_model = self.n_times
         if in_sfreq is not None and in_sfreq != preset["sfreq"]:
-            if in_sfreq != round(in_sfreq):
-                raise ValueError(
-                    f"NeuroRVQ resamples integer rates only, got {in_sfreq}."
-                )
-            warnings.warn(
-                f"NeuroRVQ-{modality.upper()} was trained at {preset['sfreq']} Hz; the "
-                f"{in_sfreq:g} Hz input is resampled inside the model.",
-                UserWarning,
-            )
-            self.resample = Resample(int(in_sfreq), preset["sfreq"])
-            # Derived from the rates: kept out of the state dict, so pretrained
-            # (200 Hz) weights load into a resampling model unchanged.
-            self.resample.register_buffer(
-                "kernel", self.resample.kernel, persistent=False
+            self.resample, self._n_times_model = _in_model_resample(
+                f"NeuroRVQ-{modality.upper()}", in_sfreq, preset["sfreq"], self.n_times
             )
-            # torchaudio's output length.
-            self._n_times_model = math.ceil(self.n_times * preset["sfreq"] / in_sfreq)
         # ceil for "pad", floor for "crop", exact for "error" (raises otherwise).
         num_patches = _patch_count(self._n_times_model, patch_size, on_non_divisible)
         if num_patches > max_patches: