braindecode.models.AttnSleep#

class braindecode.models.AttnSleep(sfreq=None, n_tce=2, d_model=80, d_ff=120, n_attn_heads=5, drop_prob=0.1, activation_mrcnn=<class 'torch.nn.modules.activation.GELU'>, activation=<class 'torch.nn.modules.activation.ReLU'>, input_window_seconds=None, n_outputs=None, after_reduced_cnn_size=30, return_feats=False, chs_info=None, n_chans=None, n_times=None)[source]#

Sleep Staging Architecture from Eldele et al. (2021) [Eldele2021].

Convolution Small Attention

AttnSleep Architecture

Attention based Neural Net for sleep staging as described in [Eldele2021]. The code for the paper and this model is also available at [1]. Takes single channel EEG as input. Feature extraction module based on multi-resolution convolutional neural network (MRCNN) and adaptive feature recalibration (AFR). The second module is the temporal context encoder (TCE) that leverages a multi-head attention mechanism to capture the temporal dependencies among the extracted features.

Warning - This model was designed for signals of 30 seconds at 100Hz or 125Hz (in which case the reference architecture from [1] which was validated on SHHS dataset [2] will be used) to use any other input is likely to make the model perform in unintended ways.

Parameters:

  • sfreq (float) – Sampling frequency of the EEG recordings.

  • n_tce (int) – Number of TCE clones.

  • d_model (int) – Input dimension for the TCE. Also the input dimension of the first FC layer in the feed forward and the output of the second FC layer in the same. Increase for higher sampling rate/signal length. It should be divisible by n_attn_heads

  • d_ff (int) – Output dimension of the first FC layer in the feed forward and the input dimension of the second FC layer in the same.

  • n_attn_heads (int) – Number of attention heads. It should be a factor of d_model

  • drop_prob (float) – Dropout rate in the PositionWiseFeedforward layer and the TCE layers.

  • activation_mrcnn (nn.Module, default=nn.ReLU) – Activation function class to apply in the Mask R-CNN layer. Should be a PyTorch activation module class like nn.ReLU or nn.GELU. Default is nn.GELU.

  • activation (nn.Module, default=nn.ReLU) – Activation function class to apply. Should be a PyTorch activation module class like nn.ReLU or nn.ELU. Default is nn.ReLU.

  • input_window_seconds (float) – Length of the input window in seconds.

  • n_outputs (int) – Number of outputs of the model. This is the number of classes in the case of classification.

  • after_reduced_cnn_size (int) – Number of output channels produced by the convolution in the AFR module.

  • return_feats (bool) – If True, return the features, i.e. the output of the feature extractor (before the final linear layer). If False, pass the features through the final linear layer.

  • chs_info (list of dict) – Information about each individual EEG channel. This should be filled with info["chs"]. Refer to mne.Info for more details.

  • n_chans (int) – Number of EEG channels.

  • n_times (int) – Number of time samples of the input window.

Raises:

ValueError – If some input signal-related parameters are not specified: and can not be inferred.

Notes

If some input signal-related parameters are not specified, there will be an attempt to infer them from the other parameters.

Hugging Face Hub integration

When the optional huggingface_hub package is installed, all models automatically gain the ability to be pushed to and loaded from the Hugging Face Hub. Install with:

pip install braindecode[hug]

Pushing a model to the Hub:

from braindecode.models import EEGNetv4

# Train your model
model = EEGNetv4(n_chans=22, n_outputs=4, n_times=1000)
# ... training code ...

# Push to the Hub
model.push_to_hub(
    repo_id="username/my-eegnet-model", commit_message="Initial model upload"
)

Loading a model from the Hub:

from braindecode.models import EEGNetv4

# Load pretrained model
model = EEGNetv4.from_pretrained("username/my-eegnet-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.

Important

Currently, only EEG-specific parameters (n_outputs, n_chans, n_times, input_window_seconds, sfreq, chs_info) are saved to the Hub. Model-specific parameters (e.g., dropout rates, activation functions, number of filters) are not preserved and will use their default values when loading from the Hub.

To use non-default model parameters, specify them explicitly when calling from_pretrained():

model = EEGNet.from_pretrained("user/model", dropout=0.3, activation='relu')

Full parameter serialization will be addressed in a future update.

References

[Eldele2021] (1,2)

E. Eldele et al., “An Attention-Based Deep Learning Approach for Sleep Stage Classification With Single-Channel EEG,” in IEEE Transactions on Neural Systems and Rehabilitation Engineering, vol. 29, pp. 809-818, 2021, doi: 10.1109/TNSRE.2021.3076234.

Methods

forward(x)[source]#

Forward pass.

Parameters:

x (torch.Tensor) – Batch of EEG windows of shape (batch_size, n_channels, n_times).

Return type:

Tensor

Examples using braindecode.models.AttnSleep#

Sleep staging on the Sleep Physionet dataset using Eldele2021

Sleep staging on the Sleep Physionet dataset using Eldele2021