EEGConformer

class selfeeg.models.zoo.EEGConformer(nb_classes: int, Chans: int, Samples: int, F: int = 40, K1: int = 25, Pool: int = 75, stride_pool: int = 15, d_model: int = 40, nlayers: int = 6, nheads: int = 10, dim_feedforward: int = 160, activation_transformer: str = 'gelu', p: float = 0.2, p_transformer: float = 0.5, mlp_dim: list[int, int] = [256, 32], return_logits: bool = True, seed: int = None)[source]

Pytorch implementation of EEGConformer.

For more information see the following paper [EEGcon] . The original implementation of EEGconformer can be found here [EEGcongit] .

The expected input is a 3D tensor with size (Batch x Channels x Samples).

Parameters:
  • nb_classes (int) – The number of classes. If less than 2, a binary classification problem is considered (output dimensions will be [batch, 1] in this case).

  • Chans (int) – The number of EEG channels.

  • Samples (int) – The sample length (number of time steps). It will be used to calculate the embedding size (for head initialization).

  • F (int, optional) –

    The number of output filters in the temporal convolution layer.

    Default = 40

  • K1 (int, optional) –

    The length of the temporal convolutional layer.

    Default = 25

  • Pool (int, optional) –

    The temporal pooling kernel size.

    Default = 75

  • stride_pool (int, optional) –

    The temporal pooling stride.

    Default = 15

  • d_model (int, optional) –

    The embedding size. It is the number of expected features in the input of the transformer encoder layer.

    Default = 40

  • nlayers (int, optional) –

    The number of transformer encoder layers.

    Default = 6

  • nheads (int, optional) –

    The number of heads in the multi-head attention layers.

    Default = 10

  • dim_feedforward (int, optional) –

    The dimension of the feedforward hidden layer in the transformer encoder.

    Default = 160

  • activation_transformer (str or Callabel, optional) –

    The activation function in the transformer encoder. See the PyTorch TransformerEncoderLayer documentation for accepted inputs.

    Default = “gelu”

  • p (float, optional) –

    Dropout probability in the tokenizer. Must be in [0,1)

    Default = 0.2

  • p_transformer (float, optional) –

    Dropout probability in the transformer encoder. Must be in [0,1)

    Default = 0.5

  • mlp_dim (list, optional) –

    A two-element list indicating the output dimensions of the 2 FC layers in the final classification head.

    Default = [256, 32]

  • return_logits (bool, optional) –

    Whether to return the output as logit or probability. It is suggested to not use False as the pytorch crossentropy loss function applies the softmax internally.

    Default = True

  • seed (int, optional) –

    A custom seed for model initialization. It must be a nonnegative number. If None is passed, no custom seed will be set

    Default = None

References

[EEGcon]

Song et al., EEG Conformer: Convolutional Transformer for EEG Decoding and Visualization. IEEE TNSRE. 2023. https://doi.org/10.1109/TNSRE.2022.3230250

Example

>>> import selfeeg.models
>>> import torch
>>> x = torch.randn(4,8,512)
>>> mdl = models.EEGConformer(2, 8, 512)
>>> out = mdl(x)
>>> print(out.shape) # shoud return torch.Size([4, 1])