Graph Neural Networks

torcheeg.models.DGCNN

class torcheeg.models.DGCNN(in_channels: int = 5, num_electrodes: int = 62, num_layers: int = 2, hid_channels: int = 32, num_classes: int = 2)[source]

Bases: Module

Dynamical Graph Convolutional Neural Networks (DGCNN). For more details, please refer to the following information.

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

dataset = SEEDDataset(io_path=f'./seed',
                      root_path='./Preprocessed_EEG',
                      offline_transform=transforms.BandDifferentialEntropy({
                          "delta": [1, 4],
                          "theta": [4, 8],
                          "alpha": [8, 14],
                          "beta": [14, 31],
                          "gamma": [31, 49]
                      }),
                      online_transform=transforms.Compose([
                          transforms.ToTensor()
                      ]),
                      label_transform=transforms.Compose([
                          transforms.Select(['emotion']),
                          transforms.Lambda(x: x + 1)
                      ]))
model = DGCNN(in_channels=5, num_electrodes=62, hid_channels=32, num_layers=2, num_classes=2)
Parameters
  • in_channels (int) – The feature dimension of each electrode. (defualt: 5)

  • num_electrodes (int) – The number of electrodes. (defualt: 62)

  • num_layers (int) – The number of graph convolutional layers. (defualt: 2)

  • hid_channels (int) – The number of hidden nodes in the first fully connected layer. (defualt: 32)

  • num_classes (int) – The number of classes to predict. (defualt: 2)

forward(x: Tensor) Tensor[source]
Parameters

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

Returns

the predicted probability that the samples belong to the classes.

Return type

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

training: bool

torcheeg.models.RGNN

class torcheeg.models.RGNN(adj: Tensor, num_electrodes: int = 62, in_channels: int = 4, num_layers: int = 2, hid_channels: int = 32, num_classes: int = 2, dropout: float = 0.7, learn_edge_weights: bool = True)[source]

Bases: Module

Regularized Graph Neural Networks (RGNN). For more details, please refer to the following information.

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

dataset = SEEDDataset(io_path=f'./seed',
                      root_path='./Preprocessed_EEG',
                      online_transform=transforms.ToG(SEED_ADJACENCY_MATRIX),,
                      label_transform=transforms.Compose([
                          transforms.Select(['emotion']),
                          transforms.Lambda(x: x + 1)
                      ]))
model = RGNN(adj=torch.Tensor(SEED_ADJACENCY_MATRIX),
             in_channels=4,
             num_electrodes=62,
             hid_channels=32,
             num_layers=2,
             num_classes=2,
             dropout=0.7,
             domain_adaptation=False,
             beta=0.0,
             learn_edge_weights=True)
Parameters
  • adj (torch.Tensor) – The adjacency matrix corresponding to the EEG representation, where 1.0 means the node is adjacent and 0.0 means the node is not adjacent. The matrix shape should be [num_electrodes, num_electrodes].

  • num_electrodes (int) – The number of electrodes. (defualt: 62)

  • in_channels (int) – The feature dimension of each electrode. (defualt: 4)

  • num_layers (int) – The number of graph convolutional layers. (defualt: 2)

  • hid_channels (int) – The number of hidden nodes in the first fully connected layer. (defualt: 32)

  • num_classes (int) – The number of classes to predict. (defualt: 2)

  • dropout (float) – Probability of an element to be zeroed in the dropout layers at the output fully-connected layer. (defualt: 0.7)

  • learn_edge_weights (bool) – Whether to learn a set of parameters to adjust the adjacency matrix. (defualt: True)

forward(data: Batch) Tensor[source]
Parameters

x (torch_geometric.data.Batch) – EEG signal representation, the ideal input shape is [n, 62, 4]. Here, n corresponds to the batch size, 62 corresponds to num_electrodes, and 4 corresponds to in_channels.

Returns

the predicted probability that the samples belong to the classes.

Return type

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

training: bool