braindecode.models.Labram#

class braindecode.models.Labram(n_times=None, n_outputs=None, chs_info=None, n_chans=None, sfreq=None, input_window_seconds=None, patch_size=200, emb_size=200, in_channels=1, out_channels=8, n_layers=12, att_num_heads=10, mlp_ratio=4.0, qkv_bias=False, qk_norm=None, qk_scale=None, drop_prob=0.0, attn_drop_prob=0.0, drop_path_prob=0.0, norm_layer=<class 'torch.nn.modules.normalization.LayerNorm'>, init_values=None, use_abs_pos_emb=True, use_mean_pooling=True, init_scale=0.001, neural_tokenizer=True, attn_head_dim=None, activation: ~torch.nn.modules.module.Module = <class 'torch.nn.modules.activation.GELU'>)[source]#

Labram from Jiang, W B et al (2024) [Jiang2024].

Labram Architecture.

Large Brain Model for Learning Generic Representations with Tremendous EEG Data in BCI from [Jiang2024]

This is an adaptation of the code [Code2024] from the Labram model.

The model is transformer architecture with strong inspiration from BEiTv2 [BeiTv2].

The models can be used in two modes: - Neural Tokenizor: Design to get an embedding layers (e.g. classification). - Neural Decoder: To extract the ampliture and phase outputs with a VQSNP.

The braindecode’s modification is to allow the model to be used in with an input shape of (batch, n_chans, n_times), if neural tokenizer equals True. The original implementation uses (batch, n_chans, n_patches, patch_size) as input with static segmentation of the input data.

The models have the following sequence of steps: if neural tokenizer:

  • SegmentPatch: Segment the input data in patches;

  • TemporalConv: Apply a temporal convolution to the segmented data;

  • Residual adding cls, temporal and position embeddings (optional);

  • WindowsAttentionBlock: Apply a windows attention block to the data;

  • LayerNorm: Apply layer normalization to the data;

  • Linear: An head linear layer to transformer the data into classes.

else:

  • PatchEmbed: Apply a patch embedding to the input data;

  • Residual adding cls, temporal and position embeddings (optional);

  • WindowsAttentionBlock: Apply a windows attention block to the data;

  • LayerNorm: Apply layer normalization to the data;

  • Linear: An head linear layer to transformer the data into classes.

Added in version 0.9.

Parameters:

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

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

  • 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.

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

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

  • patch_size (int) – The size of the patch to be used in the patch embedding.

  • emb_size (int) – The dimension of the embedding.

  • in_channels (int) – The number of convolutional input channels.

  • out_channels (int) – The number of convolutional output channels.

  • n_layers (int (default=12)) – The number of attention layers of the model.

  • att_num_heads (int (default=10)) – The number of attention heads.

  • mlp_ratio (float (default=4.0)) – The expansion ratio of the mlp layer

  • qkv_bias (bool (default=False)) – If True, add a learnable bias to the query, key, and value tensors.

  • qk_norm (Pytorch Normalize layer (default=None)) – If not None, apply LayerNorm to the query and key tensors

  • qk_scale (float (default=None)) – If not None, use this value as the scale factor. If None, use head_dim**-0.5, where head_dim = dim // num_heads.

  • drop_prob (float (default=0.0)) – Dropout rate for the attention weights.

  • attn_drop_prob (float (default=0.0)) – Dropout rate for the attention weights.

  • drop_path_prob (float (default=0.0)) – Dropout rate for the attention weights used on DropPath.

  • norm_layer (Pytorch Normalize layer (default=nn.LayerNorm)) – The normalization layer to be used.

  • init_values (float (default=None)) – If not None, use this value to initialize the gamma_1 and gamma_2 parameters.

  • use_abs_pos_emb (bool (default=True)) – If True, use absolute position embedding.

  • use_mean_pooling (bool (default=True)) – If True, use mean pooling.

  • init_scale (float (default=0.001)) – The initial scale to be used in the parameters of the model.

  • neural_tokenizer (bool (default=True)) – The model can be used in two modes: Neural Tokenizor or Neural Decoder.

  • attn_head_dim (bool (default=None)) – The head dimension to be used in the attention layer, to be used only during pre-training.

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

Raises:

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

  • FutureWarning – If add_log_softmax is True, since LogSoftmax final layer: will be removed in the future.

Notes

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

References

[Jiang2024] (1,2)

Wei-Bang Jiang, Li-Ming Zhao, Bao-Liang Lu. 2024, May. Large Brain Model for Learning Generic Representations with Tremendous EEG Data in BCI. The Twelfth International Conference on Learning Representations, ICLR.

[Code2024]

Wei-Bang Jiang, Li-Ming Zhao, Bao-Liang Lu. 2024. Labram Large Brain Model for Learning Generic Representations with Tremendous EEG Data in BCI. GitHub 935963004/LaBraM (accessed 2024-03-02)

[BeiTv2]

Zhiliang Peng, Li Dong, Hangbo Bao, Qixiang Ye, Furu Wei. 2024. BEiT v2: Masked Image Modeling with Vector-Quantized Visual Tokenizers. arXiv:2208.06366 [cs.CV]

Methods

fix_init_weight_and_init_embedding()[source]#

Fix the initial weight and the initial embedding. Initializing with truncated normal distribution.

forward(x, input_chans=None, return_patch_tokens=False, return_all_tokens=False)[source]#

Forward the input EEG data through the model.

Parameters:

  • x (torch.Tensor) – The input data with shape (batch, n_chans, n_times) or (batch, n_chans, n_patches, patch size).

  • input_chans (int) – An input channel to select some dimensions

  • return_patch_tokens (bool) – Return the patch tokens

  • return_all_tokens (bool) – Return all the tokens

Returns:

The output of the model with dimensions (batch, n_outputs)

Return type:

torch.Tensor

forward_features(x, input_chans=None, return_patch_tokens=False, return_all_tokens=False)[source]#

Forward the features of the model.

Parameters:

  • x (torch.Tensor) – The input data with shape (batch, n_chans, n_patches, patch size), if neural decoder or (batch, n_chans, n_times), if neural tokenizer.

  • input_chans (int) – The number of input channels.

  • return_patch_tokens (bool) – Whether to return the patch tokens.

  • return_all_tokens (bool) – Whether to return all the tokens.

Returns:

x – The output of the model.

Return type:

torch.Tensor

get_classifier()[source]#

Get the classifier of the model.

Returns:

The classifier of the head model.

Return type:

torch.nn.Module

get_num_layers()[source]#

Convenience method to get the number of layers in the model.

reset_classifier(n_outputs)[source]#

Reset the classifier with the new number of classes.

Parameters:

n_outputs (int) – The new number of classes.