EEGNet

class torcheeg.models.EEGNet(chunk_size: int = 151, num_electrodes: int = 60, F1: int = 8, F2: int = 16, D: int = 2, num_classes: int = 2, kernel_1: int = 64, kernel_2: int = 16, dropout: float = 0.25)

A compact convolutional neural network (EEGNet). For more details, please refer to the following information.

• Paper: Lawhern V J, Solon A J, Waytowich N R, et al. EEGNet: a compact convolutional neural network for EEG-based brain-computer interfaces[J]. Journal of neural engineering, 2018, 15(5): 056013.

• URL: https://arxiv.org/abs/1611.08024

• Related Project: https://github.com/braindecode/braindecode/tree/master/braindecode

Below is a recommended suite for use in emotion recognition tasks:

from torcheeg.datasets import DEAPDataset
from torcheeg import transforms
from torcheeg.models import EEGNet
from torch.utils.data import DataLoader

dataset = DEAPDataset(root_path='./data_preprocessed_python',
                      online_transform=transforms.Compose([
                          transforms.To2d(),
                          transforms.ToTensor(),
                      ]),
                      label_transform=transforms.Compose([
                          transforms.Select('valence'),
                          transforms.Binary(5.0),
                      ]))

model = EEGNet(chunk_size=128,
               num_electrodes=32,
               dropout=0.5,
               kernel_1=64,
               kernel_2=16,
               F1=8,
               F2=16,
               D=2,
               num_classes=2)

x, y = next(iter(DataLoader(dataset, batch_size=64)))
model(x)

Parameters:

• chunk_size (int) – Number of data points included in each EEG chunk, i.e., \(T\) in the paper. (default: 151)

• num_electrodes (int) – The number of electrodes, i.e., \(C\) in the paper. (default: 60)

• F1 (int) – The filter number of block 1, i.e., \(F_1\) in the paper. (default: 8)

• F2 (int) – The filter number of block 2, i.e., \(F_2\) in the paper. (default: 16)

• D (int) – The depth multiplier (number of spatial filters), i.e., \(D\) in the paper. (default: 2)

• num_classes (int) – The number of classes to predict, i.e., \(N\) in the paper. (default: 2)

• kernel_1 (int) – The filter size of block 1. (default: 64)

• kernel_2 (int) – The filter size of block 2. (default: 64)

• dropout (float) – Probability of an element to be zeroed in the dropout layers. (default: 0.25)

forward(x: Tensor) → Tensor

Parameters:

x (torch.Tensor) – EEG signal representation, the ideal input shape is [n, 60, 151]. Here, n corresponds to the batch size, 60 corresponds to num_electrodes, and 151 corresponds to chunk_size.

Returns:

the predicted probability that the samples belong to the classes.

Return type:

torch.Tensor[number of sample, number of classes]