braindecode 0.8.1__py3-none-any.whl → 1.1.0__py3-none-any.whl

braindecode/eegneuralnet.py

@@ -5,32 +5,36 @@
 
 
 import abc
-import logging
 import inspect
+import logging
 
 import mne
 import numpy as np
 import torch
-from skorch import NeuralNet
 from sklearn.metrics import get_scorer
+from skorch import NeuralNet
 from skorch.callbacks import BatchScoring, EpochScoring, EpochTimer, PrintLog
-from skorch.utils import noop, to_numpy, train_loss_score, valid_loss_score, is_dataset
+from skorch.helper import SliceDataset
+from skorch.utils import is_dataset, noop, to_numpy, train_loss_score, valid_loss_score
 
-from .training.scoring import (CroppedTimeSeriesEpochScoring,
-                               CroppedTrialEpochScoring, PostEpochTrainScoring)
-from .models.util import models_dict
 from .datasets.base import BaseConcatDataset, WindowsDataset
+from .models.util import models_dict
+from .training.scoring import (
+    CroppedTimeSeriesEpochScoring,
+    CroppedTrialEpochScoring,
+    PostEpochTrainScoring,
+)
 
 log = logging.getLogger(__name__)
 
 
-def _get_model(model):
-    ''' Returns the corresponding class in case the model passed is a string. '''
+def _get_model(model: str):
+    """Returns the corresponding class in case the model passed is a string."""
     if isinstance(model, str):
         if model in models_dict:
             model = models_dict[model]
         else:
-            raise ValueError(f'Unknown model name {model!r}.')
+            raise ValueError(f"Unknown model name {model!r}.")
     return model
 
 
@@ -50,7 +54,7 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
         will be left as is.
 
         """
-        kwargs = self.get_params_for('module')
+        kwargs = self.get_params_for("module")
         module = _get_model(self.module)
         module = self.initialized_instance(module, kwargs)
         # pylint: disable=attribute-defined-outside-init
@@ -61,7 +65,7 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
         # Here we parse the callbacks supplied as strings,
         # e.g. 'accuracy', to the callbacks skorch expects
         for name, cb, named_by_user in super()._yield_callbacks():
-            if name == 'str':
+            if name == "str":
                 train_cb, valid_cb = self._parse_str_callback(cb)
                 yield train_cb
                 if self.train_split is not None:
@@ -72,15 +76,13 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
     def _parse_str_callback(self, cb_supplied_name):
         scoring = get_scorer(cb_supplied_name)
         scoring_name = scoring._score_func.__name__
-        assert scoring_name.endswith(
-            ('_score', '_error', '_deviance', '_loss'))
-        if (scoring_name.endswith('_score') or
-                cb_supplied_name.startswith('neg_')):
+        assert scoring_name.endswith(("_score", "_error", "_deviance", "_loss"))
+        if scoring_name.endswith("_score") or cb_supplied_name.startswith("neg_"):
             lower_is_better = False
         else:
             lower_is_better = True
-        train_name = f'train_{cb_supplied_name}'
-        valid_name = f'valid_{cb_supplied_name}'
+        train_name = f"train_{cb_supplied_name}"
+        valid_name = f"valid_{cb_supplied_name}"
         if self.cropped:
             train_scoring = CroppedTrialEpochScoring(
                 cb_supplied_name, lower_is_better, on_train=True, name=train_name
@@ -98,7 +100,7 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
         named_by_user = True
         train_valid_callbacks = [
             (train_name, train_scoring, named_by_user),
-            (valid_name, valid_scoring, named_by_user)
+            (valid_name, valid_scoring, named_by_user),
         ]
         return train_valid_callbacks
 
@@ -108,8 +110,13 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
         if not training:
             epoch_cbs = []
             for name, cb in self.callbacks_:
-                if isinstance(cb, (CroppedTrialEpochScoring, CroppedTimeSeriesEpochScoring)) and (
-                        hasattr(cb, 'window_inds_')) and (not cb.on_train):
+                if (
+                    isinstance(
+                        cb, (CroppedTrialEpochScoring, CroppedTimeSeriesEpochScoring)
+                    )
+                    and (hasattr(cb, "window_inds_"))
+                    and (not cb.on_train)
+                ):
                     epoch_cbs.append(cb)
             # for trialwise decoding stuffs it might also be we don't have
             # cropped loader, so no indices there
@@ -136,8 +143,11 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
         i_window_stops = np.concatenate(i_window_stops)
         window_ys = np.concatenate(window_ys)
         return dict(
-            preds=preds, i_window_in_trials=i_window_in_trials,
-            i_window_stops=i_window_stops, window_ys=window_ys)
+            preds=preds,
+            i_window_in_trials=i_window_in_trials,
+            i_window_stops=i_window_stops,
+            window_ys=window_ys,
+        )
 
     # Changes the default target extractor to noop
     @property
@@ -156,7 +166,9 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
             (
                 "valid_loss",
                 BatchScoring(
-                    valid_loss_score, name="valid_loss", target_extractor=noop,
+                    valid_loss_score,
+                    name="valid_loss",
+                    target_extractor=noop,
                 ),
             ),
             ("print_log", PrintLog()),
@@ -179,17 +191,27 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
             return
         # get kwargs from signal:
         signal_kwargs = dict()
-        if isinstance(X, mne.BaseEpochs) or isinstance(X, np.ndarray):
+        # Using shape to work both with torch.tensor and numpy.array:
+        if (
+            isinstance(X, mne.BaseEpochs)
+            or (hasattr(X, "shape") and len(X.shape) >= 2)
+            or isinstance(X, SliceDataset)
+        ):
             if y is None:
-                raise ValueError("y must be specified if X is a numpy array.")
-            signal_kwargs['n_outputs'] = self._get_n_outputs(y, classes)
+                raise ValueError("y must be specified if X is array-like.")
+            signal_kwargs["n_outputs"] = self._get_n_outputs(y, classes)
             if isinstance(X, mne.BaseEpochs):
                 self.log.info("Using mne.Epochs to find signal-related parameters.")
                 signal_kwargs["n_times"] = len(X.times)
-                signal_kwargs["sfreq"] = X.info['sfreq']
-                signal_kwargs["chs_info"] = X.info['chs']
+                signal_kwargs["sfreq"] = X.info["sfreq"]
+                signal_kwargs["chs_info"] = X.info["chs"]
+            elif isinstance(X, SliceDataset):
+                self.log.info("Using SliceDataset to find signal-related parameters.")
+                Xshape = X[0].shape
+                signal_kwargs["n_times"] = Xshape[-1]
+                signal_kwargs["n_chans"] = Xshape[-2]
             else:
-                self.log.info("Using numpy array to find signal-related parameters.")
+                self.log.info("Using array-like to find signal-related parameters.")
                 signal_kwargs["n_times"] = X.shape[-1]
                 signal_kwargs["n_chans"] = X.shape[-2]
         elif is_dataset(X):
@@ -198,21 +220,17 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
             Xshape = X0.shape
             signal_kwargs["n_times"] = Xshape[-1]
             signal_kwargs["n_chans"] = Xshape[-2]
-            if (
-                    isinstance(X, BaseConcatDataset) and
-                    all(ds.targets_from == 'metadata' for ds in X.datasets)
+            if isinstance(X, BaseConcatDataset) and all(
+                ds.targets_from == "metadata" for ds in X.datasets
             ):
                 y_target = X.get_metadata().target
-                signal_kwargs['n_outputs'] = self._get_n_outputs(y_target, classes)
-            elif (
-                    isinstance(X, WindowsDataset) and
-                    X.targets_from == "metadata"
-            ):
+                signal_kwargs["n_outputs"] = self._get_n_outputs(y_target, classes)
+            elif isinstance(X, WindowsDataset) and X.targets_from == "metadata":
                 y_target = X.windows.metadata.target
-                signal_kwargs['n_outputs'] = self._get_n_outputs(y_target, classes)
+                signal_kwargs["n_outputs"] = self._get_n_outputs(y_target, classes)
         else:
             self.log.warning(
-                "Can only infer signal shape of numpy arrays or and Datasets, "
+                "Can only infer signal shape of array-like and Datasets, "
                 f"got {type(X)!r}."
             )
             return
@@ -227,15 +245,13 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
             if k in all_module_kwargs:
                 module_kwargs[k] = v
             else:
-                self.log.warning(
-                    f"Module {self.module!r} "
-                    f"is missing parameter {k!r}."
-                )
+                self.log.warning(f"Module {self.module!r} is missing parameter {k!r}.")
 
         # save kwargs to self:
         self.log.info(
             f"Passing additional parameters {module_kwargs!r} "
-            f"to module {self.module!r}.")
+            f"to module {self.module!r}."
+        )
         module_kwargs = {f"module__{k}": v for k, v in module_kwargs.items()}
         self.set_params(**module_kwargs)
 
@@ -275,7 +291,7 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
 
         """
         if isinstance(X, mne.BaseEpochs):
-            X = X.get_data(units='uV')
+            X = X.get_data(units="uV")
         return super().get_dataset(X, y)
 
     def partial_fit(self, X, y=None, classes=None, **fit_params):
@@ -291,7 +307,7 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
 
           * mne.Epochs: ``n_times``, ``n_chans``, ``n_outputs``, ``chs_info``,
             ``sfreq``, ``input_window_seconds``
-          * numpy array: ``n_times``, ``n_chans``, ``n_outputs``
+          * array-like: ``n_times``, ``n_chans``, ``n_outputs``
           * WindowsDataset with ``targets_from='metadata'``
             (or BaseConcatDataset of such datasets): ``n_times``, ``n_chans``, ``n_outputs``
           * other Dataset: ``n_times``, ``n_chans``
@@ -345,7 +361,7 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
 
           * mne.Epochs: ``n_times``, ``n_chans``, ``n_outputs``, ``chs_info``,
             ``sfreq``, ``input_window_seconds``
-          * numpy array: ``n_times``, ``n_chans``, ``n_outputs``
+          * array-like: ``n_times``, ``n_chans``, ``n_outputs``
           * WindowsDataset with ``targets_from='metadata'``
             (or BaseConcatDataset of such datasets): ``n_times``, ``n_chans``, ``n_outputs``
           * other Dataset: ``n_times``, ``n_chans``

braindecode/functional/__init__.py

@@ -0,0 +1,22 @@
+from .functions import (
+    _get_gaussian_kernel1d,
+    drop_path,
+    hilbert_freq,
+    identity,
+    plv_time,
+    safe_log,
+    square,
+)
+from .initialization import glorot_weight_zero_bias, rescale_parameter
+
+__all__ = [
+    "_get_gaussian_kernel1d",
+    "drop_path",
+    "hilbert_freq",
+    "identity",
+    "plv_time",
+    "safe_log",
+    "square",
+    "glorot_weight_zero_bias",
+    "rescale_parameter",
+]

braindecode/functional/functions.py

@@ -0,0 +1,250 @@
+# Authors: Robin Schirrmeister <robintibor@gmail.com>
+#
+# License: BSD (3-clause)
+
+import torch
+import torch.nn.functional as F
+
+
+def square(x):
+    return x * x
+
+
+def safe_log(x, eps: float = 1e-6) -> torch.Tensor:
+    """Prevents :math:`log(0)` by using :math:`log(max(x, eps))`."""
+    return torch.log(torch.clamp(x, min=eps))
+
+
+def identity(x):
+    return x
+
+
+def drop_path(
+    x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
+):
+    """Drop paths (Stochastic Depth) per sample.
+
+
+    Notes: This implementation is taken from timm library.
+
+    All credit goes to Ross Wightman.
+
+    Parameters
+    ----------
+    x: torch.Tensor
+        input tensor
+    drop_prob : float, optional
+        survival rate (i.e. probability of being kept), by default 0.0
+    training : bool, optional
+        whether the model is in training mode, by default False
+    scale_by_keep : bool, optional
+        whether to scale output by (1/keep_prob) during training, by default True
+
+    Returns
+    -------
+    torch.Tensor
+        output tensor
+
+    Notes from Ross Wightman:
+    (when applied in main path of residual blocks)
+    This is the same as the DropConnect impl I created for EfficientNet,
+    etc. networks, however,
+    the original name is misleading as 'Drop Connect' is a different form
+    of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956
+    ... I've opted for changing the layer and argument names to 'drop path'
+    rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+
+    """
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (
+        x.ndim - 1
+    )  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+
+def _get_gaussian_kernel1d(kernel_size: int, sigma: float) -> torch.Tensor:
+    """
+    Generates a 1-dimensional Gaussian kernel based on the specified kernel
+    size and standard deviation (sigma).
+    This kernel is useful for Gaussian smoothing or filtering operations in
+    image processing. The function calculates a range limit to ensure the kernel
+    effectively covers the Gaussian distribution. It generates a tensor of
+    specified size and type, filled with values distributed according to a
+    Gaussian curve, normalized using a softmax function
+    to ensure all weights sum to 1.
+
+
+    Parameters
+    ----------
+    kernel_size: int
+    sigma: float
+
+    Returns
+    -------
+    kernel1d: torch.Tensor
+
+    Notes
+    -----
+    Code copied and modified from TorchVision:
+    https://github.com/pytorch/vision/blob/main/torchvision/transforms/_functional_tensor.py#L725-L732
+    All rights reserved.
+
+    LICENSE in https://github.com/pytorch/vision/blob/main/LICENSE
+
+    """
+    ksize_half = (kernel_size - 1) * 0.5
+    x = torch.linspace(-ksize_half, ksize_half, steps=kernel_size)
+    pdf = torch.exp(-0.5 * (x / sigma).pow(2))
+    kernel1d = pdf / pdf.sum()
+    return kernel1d
+
+
+def hilbert_freq(x, forward_fourier=True):
+    r"""
+    Compute the Hilbert transform using PyTorch, separating the real and
+    imaginary parts.
+
+    The analytic signal :math:`x_a(t)` of a real-valued signal :math:`x(t)`
+    is defined as:
+
+    .. math::
+
+        x_a(t) = x(t) + i y(t) = \mathcal{F}^{-1} \{ U(f) \mathcal{F}\{x(t)\} \}
+
+    where:
+    - :math:`\mathcal{F}` is the Fourier transform,
+    - :math:`U(f)` is the unit step function,
+    - :math:`y(t)` is the Hilbert transform of :math:`x(t)`.
+
+
+    Parameters
+    ----------
+    input : torch.Tensor
+        Input tensor. The expected shape depends on the `forward_fourier` parameter:
+
+        - If `forward_fourier` is True:
+            (..., seq_len)
+        - If `forward_fourier` is False:
+            (..., seq_len / 2 + 1, 2)
+
+    forward_fourier : bool, optional
+        Determines the format of the input tensor.
+        - If True, the input is in the forward Fourier domain.
+        - If False, the input contains separate real and imaginary parts.
+        Default is True.
+
+    Returns
+    -------
+    torch.Tensor
+        Output tensor with shape (..., seq_len, 2), where the last dimension represents
+        the real and imaginary parts of the Hilbert transform.
+
+    Examples
+    --------
+    >>> import torch
+    >>> input = torch.randn(10, 100)  # Example input tensor
+    >>> output = hilbert_transform(input)
+    >>> print(output.shape)
+    torch.Size([10, 100, 2])
+
+    Notes
+    -----
+    The implementation is matching scipy implementation, but using torch.
+    https://github.com/scipy/scipy/blob/v1.14.1/scipy/signal/_signaltools.py#L2287-L2394
+
+    """
+    if forward_fourier:
+        x = torch.fft.rfft(x, norm=None, dim=-1)
+        x = torch.view_as_real(x)
+    x = x * 2.0
+    x[..., 0, :] = x[..., 0, :] / 2.0  # Don't multiply the DC-term by 2
+    x = F.pad(
+        x, [0, 0, 0, x.shape[-2] - 2]
+    )  # Fill Fourier coefficients to retain shape
+    x = torch.view_as_complex(x)
+    x = torch.fft.ifft(x, norm=None, dim=-1)  # returns complex signal
+    x = torch.view_as_real(x)
+
+    return x
+
+
+def plv_time(x, forward_fourier=True, epsilon: float = 1e-6):
+    """Compute the Phase Locking Value (PLV) metric in the time domain.
+
+    The Phase Locking Value (PLV) is a measure of the synchronization between
+    different channels by evaluating the consistency of phase differences
+    over time. It ranges from 0 (no synchronization) to 1 (perfect
+    synchronization) [1]_.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Input tensor containing the signal data.
+        - If `forward_fourier` is `True`, the shape should be `(..., channels, time)`.
+        - If `forward_fourier` is `False`, the shape should be `(..., channels, freqs, 2)`,
+          where the last dimension represents the real and imaginary parts.
+    forward_fourier : bool, optional
+        Specifies the format of the input tensor `x`.
+        - If `True`, `x` is assumed to be in the time domain.
+        - If `False`, `x` is assumed to be in the Fourier domain with separate real and
+          imaginary components.
+        Default is `True`.
+    epsilon : float, default 1e-6
+        Small numerical value to ensure positivity constraint on the complex part
+
+    Returns
+    -------
+    plv : torch.Tensor
+        The Phase Locking Value matrix with shape `(..., channels, channels)`. Each
+        element `[i, j]` represents the PLV between channel `i` and channel `j`.
+
+    References
+    ----------
+    [1] Lachaux, J. P., Rodriguez, E., Martinerie, J., & Varela, F. J. (1999).
+        Measuring phase synchrony in brain signals. Human brain mapping,
+        8(4), 194-208.
+    """
+    # Compute the analytic signal using the Hilbert transform.
+    # x_a has separate real and imaginary parts.
+    analytic_signal = hilbert_freq(x, forward_fourier)
+    # Calculate the amplitude (magnitude) of the analytic signal.
+    # Adding a small epsilon (1e-6) to avoid division by zero.
+    amplitude = torch.sqrt(
+        analytic_signal[..., 0] ** 2 + analytic_signal[..., 1] ** 2 + 1e-6
+    )
+    # Normalize the analytic signal to obtain unit vectors (phasors).
+    unit_phasor = analytic_signal / amplitude.unsqueeze(-1)
+
+    # Compute the real part of the outer product between phasors of
+    # different channels.
+    real_real = torch.matmul(unit_phasor[..., 0], unit_phasor[..., 0].transpose(-2, -1))
+
+    # Compute the imaginary part of the outer product between phasors of
+    # different channels.
+    imag_imag = torch.matmul(unit_phasor[..., 1], unit_phasor[..., 1].transpose(-2, -1))
+
+    # Compute the cross-terms for the real and imaginary parts.
+    real_imag = torch.matmul(unit_phasor[..., 0], unit_phasor[..., 1].transpose(-2, -1))
+    imag_real = torch.matmul(unit_phasor[..., 1], unit_phasor[..., 0].transpose(-2, -1))
+
+    # Combine the real and imaginary parts to form the complex correlation.
+    correlation_real = real_real + imag_imag
+    correlation_imag = real_imag - imag_real
+
+    # Determine the number of time points (or frequency bins if in Fourier domain).
+    time = amplitude.shape[-1]
+
+    # Calculate the PLV by averaging the magnitude of the complex correlation over time.
+    # epsilon is small numerical value to ensure positivity constraint on the complex part
+    plv_matrix = (
+        1 / time * torch.sqrt(correlation_real**2 + correlation_imag**2 + epsilon)
+    )
+
+    return plv_matrix

braindecode/functional/initialization.py

@@ -0,0 +1,47 @@
+import math
+
+from torch import nn
+
+
+def glorot_weight_zero_bias(model):
+    """Initialize parameters of all modules by initializing weights with
+    glorot
+     uniform/xavier initialization, and setting biases to zero. Weights from
+     batch norm layers are set to 1.
+
+    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)
+
+
+def rescale_parameter(param, layer_id):
+    r"""Recaling the l-th transformer layer.
+
+    Rescales the parameter tensor by the inverse square root of the layer id.
+    Made inplace. :math:`\frac{1}{\sqrt{2 \cdot \text{layer\_id}}}` [Beit2022]
+
+    In the labram, this is used to rescale the output matrices
+    (i.e., the last linear projection within each sub-layer) of the
+    self-attention module.
+
+    Parameters
+    ----------
+    param: :class:`torch.Tensor`
+        tensor to be rescaled
+    layer_id: int
+        layer id in the neural network
+
+    References
+    ----------
+    [Beit2022] Hangbo Bao, Li Dong, Songhao Piao, Furu We (2022). BEIT: BERT
+    Pre-Training of Image Transformers.
+    """
+    param.div_(math.sqrt(2.0 * layer_id))

braindecode/models/__init__.py

@@ -1,30 +1,100 @@
 """
 Some predefined network architectures for EEG decoding.
 """
+
+from .atcnet import ATCNet
+from .attentionbasenet import AttentionBaseNet
 from .base import EEGModuleMixin
-from .eegconformer import EEGConformer
-from .eegitnet import EEGITNet
+from .biot import BIOT
+from .contrawr import ContraWR
+from .ctnet import CTNet
 from .deep4 import Deep4Net
 from .deepsleepnet import DeepSleepNet
-from .eegnet import EEGNetv4, EEGNetv1
-from .hybrid import HybridNet
-from .shallow_fbcsp import ShallowFBCSPNet
-from .eegresnet import EEGResNet
-from .eeginception import EEGInception
+from .eegconformer import EEGConformer
 from .eeginception_erp import EEGInceptionERP
 from .eeginception_mi import EEGInceptionMI
-from .atcnet import ATCNet
-from .tcn import TCN
-from .sleep_stager_chambon_2018 import SleepStagerChambon2018
+from .eegitnet import EEGITNet
+from .eegminer import EEGMiner
+from .eegnet import EEGNetv1, EEGNetv4
+from .eegnex import EEGNeX
+from .eegresnet import EEGResNet
+from .eegsimpleconv import EEGSimpleConv
+from .eegtcnet import EEGTCNet
+from .fbcnet import FBCNet
+from .fblightconvnet import FBLightConvNet
+from .fbmsnet import FBMSNet
+from .hybrid import HybridNet
+from .ifnet import IFNet
+from .labram import Labram
+from .msvtnet import MSVTNet
+from .sccnet import SCCNet
+from .shallow_fbcsp import ShallowFBCSPNet
+from .signal_jepa import (
+    SignalJEPA,
+    SignalJEPA_Contextual,
+    SignalJEPA_PostLocal,
+    SignalJEPA_PreLocal,
+)
+from .sinc_shallow import SincShallowNet
 from .sleep_stager_blanco_2020 import SleepStagerBlanco2020
+from .sleep_stager_chambon_2018 import SleepStagerChambon2018
 from .sleep_stager_eldele_2021 import SleepStagerEldele2021
+from .sparcnet import SPARCNet
+from .syncnet import SyncNet
+from .tcn import BDTCN, TCN
 from .tidnet import TIDNet
+from .tsinception import TSceptionV1
 from .usleep import USleep
-from .util import get_output_shape, to_dense_prediction_model
-from .modules import TimeDistributed
-
-from .util import _init_models_dict
+from .util import _init_models_dict, models_mandatory_parameters
 
 # Call this last in order to make sure the dataset list is populated with
 # the models imported in this file.
 _init_models_dict()
+
+__all__ = [
+    "ATCNet",
+    "AttentionBaseNet",
+    "EEGModuleMixin",
+    "BIOT",
+    "ContraWR",
+    "CTNet",
+    "Deep4Net",
+    "DeepSleepNet",
+    "EEGConformer",
+    "EEGInceptionERP",
+    "EEGInceptionMI",
+    "EEGITNet",
+    "EEGMiner",
+    "EEGNetv1",
+    "EEGNetv4",
+    "EEGNeX",
+    "EEGResNet",
+    "EEGSimpleConv",
+    "EEGTCNet",
+    "FBCNet",
+    "FBLightConvNet",
+    "FBMSNet",
+    "HybridNet",
+    "IFNet",
+    "Labram",
+    "MSVTNet",
+    "SCCNet",
+    "ShallowFBCSPNet",
+    "SignalJEPA",
+    "SignalJEPA_Contextual",
+    "SignalJEPA_PostLocal",
+    "SignalJEPA_PreLocal",
+    "SincShallowNet",
+    "SleepStagerBlanco2020",
+    "SleepStagerChambon2018",
+    "SleepStagerEldele2021",
+    "SPARCNet",
+    "SyncNet",
+    "BDTCN",
+    "TCN",
+    "TIDNet",
+    "TSceptionV1",
+    "USleep",
+    "_init_models_dict",
+    "models_mandatory_parameters",
+]