Training models with PyTorch Geometric

In this case, we introduce how to use TorchEEG and a customized training process based on vanilla PyTorch to train a PyTorch Geometric-based graph convolutional network on the SEED dataset for emotion classification.

import logging
import os
import random
import time

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import Linear
from torch_geometric.loader import DataLoader
from torch_geometric.nn import GATConv, global_mean_pool

from torcheeg import transforms
from torcheeg.datasets import SEEDFeatureDataset
from torcheeg.datasets.constants.emotion_recognition.seed import \
    SEED_ADJACENCY_MATRIX
from torcheeg.model_selection import KFoldPerSubjectGroupbyTrial
from torcheeg.transforms.pyg import ToG

Pre-experiment Preparation to Ensure Reproducibility

Use the logging module to store output in a log file for easy reference while printing it to the screen.

os.makedirs('./tmp_out/examples_torch/log', exist_ok=True)
logger = logging.getLogger('Training models with vanilla PyTorch')
logger.setLevel(logging.DEBUG)
console_handler = logging.StreamHandler()
timeticks = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
file_handler = logging.FileHandler(
    os.path.join('./tmp_out/examples_torch/log', f'{timeticks}.log'))
logger.addHandler(console_handler)
logger.addHandler(file_handler)

Set the random number seed in all modules to guarantee the same result when running again.

def seed_everything(seed):
    random.seed(seed)
    np.random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


seed_everything(42)

Defining a graph convolutional network

Use the API provided by PyTorch Geometric to define graph convolutional networks. Here, the EEG signal or feature of the electrode corresponds to the input :obj`data.x`, and the relationship between electrodes corresponds to :obj`data.edge_index`. Depending on the definition of the adjacency matrix, the relationship may represent spatial adjacency, etc.

class GNN(torch.nn.Module):
    def __init__(self,
                 in_channels=5,
                 num_layers=3,
                 hid_channels=64,
                 num_classes=3):
        super().__init__()
        self.conv1 = GATConv(in_channels, hid_channels)
        self.convs = torch.nn.ModuleList()
        for _ in range(num_layers - 1):
            self.convs.append(GATConv(hid_channels, hid_channels))
        self.lin1 = Linear(hid_channels, hid_channels)
        self.lin2 = Linear(hid_channels, num_classes)

    def reset_parameters(self):
        self.conv1.reset_parameters()
        for conv in self.convs:
            conv.reset_parameters()
        self.lin1.reset_parameters()
        self.lin2.reset_parameters()

    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch
        x = F.relu(self.conv1(x, edge_index))
        for conv in self.convs:
            x = F.relu(conv(x, edge_index))
        x = global_mean_pool(x, batch)
        x = F.relu(self.lin1(x))
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.lin2(x)
        return x

Customize the training process

TorchEEG provides a large number of trainers to help complete the training of classification models, however, you can also define the training functions to complete the training and testing of the model. Here is a simple example:

# training process
def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    model.train()
    for batch_idx, batch in enumerate(dataloader):
        X = batch[0].to(device)
        y = batch[1].to(device)

        # Compute prediction error
        pred = model(X)
        loss = loss_fn(pred, y)

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch_idx % 20 == 0:
            loss, current = loss.item(), batch_idx * batch_size
            logger.info(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")


# validation process
def valid(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    val_loss, correct = 0, 0
    with torch.no_grad():
        for batch in dataloader:
            X = batch[0].to(device)
            y = batch[1].to(device)

            pred = model(X)
            val_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    val_loss /= num_batches
    correct /= size
    logger.info(
        f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {val_loss:>8f} \n"
    )

Building Deep Learning Pipelines Using TorchEEG

Step 1: Initialize the Dataset

We use the SEED dataset supported by TorchEEG. Here we are using extracted EEG features. In the feature dataset, EEG signals (200 data points) per second are pre-computed with differential entropy in five sub-bands and smoothed using a linear dynamical system. .. note:

In online processing, All EEG signals are converted into a graph structure that can be processed by the PyTorch Geometric model, i.e., :obj:`torch_geometric.data.Data`, according to the adjacency matrix and the signals or features from electrodes. Here, electrodes represent nodes, and the adjacency matrix defines the association between electrodes.

dataset = SEEDFeatureDataset(io_path='./tmp_out/examples_torch_geometric/seed',
                             root_path='./tmp_in/ExtractedFeatures',
                             feature=['de_movingAve'],
                             online_transform=transforms.Compose([
                                 transforms.MinMaxNormalize(axis=-1),
                                 ToG(SEED_ADJACENCY_MATRIX)
                             ]),
                             label_transform=transforms.Compose([
                                 transforms.Select('emotion'),
                                 transforms.Lambda(lambda x: int(x) + 1),
                             ]),
                             num_worker=8)

Warning

If you use TorchEEG under the Windows system and want to use multiple processes (such as in dataset or dataloader), you should check whether __name__ is __main__ to avoid errors caused by multiple import.

That is, under the Windows system, you need to:

if __name__ == "__main__":
    dataset = SEEDFeatureDataset(io_path='./tmp_out/examples_torch_geometric/seed',
                          root_path='./tmp_in/ExtractedFeatures',
                          feature=['de_movingAve'],
                          online_transform=transforms.Compose([
                              transforms.MinMaxNormalize(axis=-1),
                              ToG(SEED_ADJACENCY_MATRIX)
                          ]),
                          label_transform=transforms.Compose([
                              transforms.Select('emotion'),
                              transforms.Lambda(lambda x: int(x) + 1),
                          ]),
                          io_mode='pickle',
                          num_worker=8)
    # the following codes

Note

LMDB may not be optimized for parts of Windows systems or storage devices. If you find that the data preprocessing speed is slow, you can consider setting io_mode to pickle, which is an alternative implemented by TorchEEG based on pickle.

Step 2: Divide the Training and Test samples in the Dataset

Here, the dataset is divided using per-subject 5-fold cross-validation. In the process of division, we split the training and test sets separately on each subject’s EEG samples. Here, we take 4 folds as training samples and 1 fold as test samples.

k_fold = KFoldPerSubjectGroupbyTrial(
    n_splits=10,
    split_path='./tmp_out/examples_torch_geometric/split',
    shuffle=False)

Step 3: Define the Model and Start Training

We first use a loop to get the dataset in each cross-validation. In each cross-validation, we initialize the above-mentioned GNN model.

Next, we train the model for 50 epochs using the training function defined above and report the model performance on the validation set at each epoch with the validation function defined above.

Note

Please note that since the EEG signal sample returned by ToG is of type torch_geometric.data.Data, which represents a graph structure, the DataLoader provided by PyTorch cannot be used here to form a batch. Instead, torch_geometric.loader.DataLoader should be used to batch the adjacency matrix of the graph structure and the node features.

device = "cuda" if torch.cuda.is_available() else "cpu"
loss_fn = nn.CrossEntropyLoss()
batch_size = 64

for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
    # initialize model
    model = GNN().to(device)
    # initialize optimizer
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

    train_loader = DataLoader(train_dataset,
                              batch_size=batch_size,
                              shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

    epochs = 50
    for t in range(epochs):
        logger.info(f"Epoch {t+1}\n-------------------------------")
        train(train_loader, model, loss_fn, optimizer)
        valid(val_loader, model, loss_fn)