ATCNet

class selfeeg.models.zoo.ATCNet(nb_classes: int, Chans: int, Samples: int, Fs: float, num_windows: int = 5, mha_heads: int = 2, tcn_depth: int = 2, F1: int = 16, D: int = 2, pool1: int = None, pool2: int = None, dropRate: float = 0.3, max_norm: float = None, batchMomentum: float = 0.1, ELUAlpha: float = 1.0, mha_dropRate: float = 0.5, tcn_kernLength: int = 4, tcn_F: int = 32, tcn_ELUAlpha: float = 0.0, tcn_dropRate: float = 0.3, tcn_max_norm: float = None, tcn_batchMom: float = 0.1, return_logits: bool = True, seed: int = None)[source]

Pytorch implementation of the ATCNet model.

ATCNet paper can be found here [atcnet] . The official implementation can be found here [gitatc] .

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

Warning

Due to the multi-branch nature of the network and the usage of a classification head at the end of each branch, this model does not have an implementation of the encoder. Keep in mind that the first convolutional block is basically an eegnet encoder with a conv2d instead of a separable conv2d.

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 number of EEG samples.

  • Fs (int or float) – The EEG sampling rate.

  • num_windows (int, optonal) –

    The number of branches to use after the first convolutional block.

    Default = 5

  • mha_heads (int, optional) –

    The number of multi-head attention heads to use in each branch.

    Default = 2

  • tcn_depth (int, optional) –

    The number of temporal convolution blocks to use in each branch.

    Default = 2

  • F1 (int, optional) –

    The number of filters to use in the first layer of the first convolutional block. It is the same as the F1 argument of the EEGNet model

    Default = 16

  • D (int, optional) –

    The depth of the depthwise layer in the first convolutional block. It is the same as the D argument of the EEGNet model.

    Default = 2

  • pool1 (int, optional) –

    The kernel length of the first pooling layer of the first convolutional block. It is the same as the pool1 argument of the EEGNet model. If left to None the length is automatically retrieved to cover the same portion of EEG as in the original work (regardless of the Sampling Rate).

    Default = None

  • pool2 (int, optional) –

    The kernel length of the second pooling layer of the first convolutional block. It is the same as the pool2 argument of the EEGNet model. If left to None the length is automatically retrieved to cover the same portion of EEG as in the original work (regardless of the Sampling Rate).

    Default = None

  • dropRate (float, optional) –

    The dropout rate of the first convolutional layer. It is the same as the dropRate argument in the EEGNet model.

    Default = 0.3

  • max_norm (float, optional) –

    The max norm constraint to apply to the convolutional layers of the first convolutioal block. If left to None, no constraints will be applied

    Default = None

  • batchMomentum (float, optional) –

    The batch normalization momentum of the first convolutional layer. It is the same as the batch_momentum argument of the EEGNet model. Note that the original paper uses a higher batch momentum (0.9).

    Default = 0.1

  • ELUAlpha (float, optional) –

    the alpha value of the ELU activation function. It is the same as the batch_momentum argument of the EEGNet model.

    Default = 1

  • mha_dropRate (float, optional) –

    The dropout rate of the multi head attention block.

    Default = 0.5

  • tcn_kernLength (int, optional) –

    The kernel length of the temporal convolutional block.

    Default = 4

  • tcn_F (int, optional) –

    The number of filters of the temporal convolutioal block.

    Default = 32

  • tcn_ELUAlpha (float, optional) –

    The alpha value for the activation function of the temporal convolutioal block.

    Default = 1.0

  • tcn_dropRate (float, optional) –

    The dropout rate of the temporal convolutioal block.

    Default = 0.5

  • tcn_max_norm (float, optional) – The max norm constraint to apply to the convolutional layer of the temporal convolutioal block. If left to None, no constraints will be applied

  • tcn_batchMom (float, optional) – The batch normalization momentum of the temporal convolutioal block.

  • return_logits (bool, optional) –

    Whether to return the output as logit or probability. It is suggested to not use False as the pytorch crossentropy 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

[atcnet]

Altaheri et al. Physics-Informed Attention Temporal Convolutional Network for EEG-Based Motor Imagery Classification. https://ieeexplore.ieee.org/document/9852687

Example

>>> import selfeeg.models
>>> import torch
>>> x = torch.randn(4,8,512)
>>> mdl = models.ATCNet(2, 8, 512, 128)
>>> out = mdl(x)
>>> print(out.shape) # shoud return torch.Size([4, 2])
>>> print(torch.isnan(out).sum()) # shoud return 0