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.
Paper: Song T, Zheng W, Song P, et al. EEG emotion recognition using dynamical graph convolutional neural networks[J]. IEEE Transactions on Affective Computing, 2018, 11(3): 532-541.
Related Project: https://github.com/xueyunlong12589/DGCNN
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 tonum_electrodes
, and5
corresponds toin_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.
Paper: Zhong P, Wang D, Miao C. EEG-based emotion recognition using regularized graph neural networks[J]. IEEE Transactions on Affective Computing, 2020.
Related Project: https://github.com/zhongpeixiang/RGNN
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 tonum_electrodes
, and4
corresponds toin_channels
.- Returns
the predicted probability that the samples belong to the classes.
- Return type
torch.Tensor[number of sample, number of classes]
- training: bool