Spatio-Temporal EEG Representation Learning on Riemannian Manifold and Euclidean Space

pip install eeg-riemannian==1.10.0

This repository contains the source code of our paper, using following datasets:

• Emotion Recoginition:

  • SEED: 15 subjects participated experiments with videos as emotion stimuli (positive/negative/neutral) and EEG was recorded with 62 channels at sampling rate of 1000Hz.

  • SEED-VIG: Vigilance estimation using EEG data in a simulated driving task. 23 subjects participated experiments and 17 EEG channels were recorded at sampling rate of 1000Hz.

• Motor Imagery:

  • BCI-IV 2a: 9 subjects were involved in motor-imagery experiment (left hand, right hand, feet and tongue). 22 EEG recordings were collected at sampling rate of 250Hz.

  • BCI-IV 2b: 9 subjects were involved in motor-imagery experiment (left hand and right hand). 3 EEG channels were recorded at sampling rate of 250Hz.

Please follow the steps below in order to be able to train our models:

1 - Install Requirements

pip3 install -r ./requirements.txt

2 - Download dataset, load data, preprocess data through filter bank, then perform feature extraction.

3 - Save the preprocessed data and EEG into separate folders (e.g., '/train/EEG/' and '/train/Extracted Features'). Move data and corresponding labels to the address shown in functions 'load_dataset_signal_addr' and 'load_dataset_feature_addr' from utils.

4 - Perform a hyper-parameters search for each information stream.

(1) For spatial information stream, run python3 ./main_spatial_val.py --dataset datasetname to search for the rank of EEG covariance matrices. For example, run the following command

python3 ./main_spatial_val.py --dataset BCI_IV_2a --cpu-seed 0 --gpu-seed 12345 --lr 0.001 --batch-size 32 --epochs 200 --early-stopping 20 --riemannian_dist

for the BCI_IV_2a dataset using Riemannian projection.

(2) For temporal information stream, run python3 ./main_temporal_val.py --dataset datasetname to obtain the result for different LSTM settings. For example, run the following command

python3 ./main_temporal_val.py --dataset SEED --cpu-seed 0 --gpu-seed 12345 --lr 0.001 --batch-size 8 --epochs 200 --early-stopping 20 -- BiLSTM --layer-num 2

for the SEED dataset using two bidirectional LSTM layers.

Validation results will be automatically saved in the address in functions 'save_spatial_val_result' and 'save_temporal_val_result' from utils. The parameters are saved and updated in dataset_params.

5 - Run the experiments for test data. For example, run the following command

python3 ./main.py --dataset BCI_IV_2b --cpu-seed 0 --gpu-seed 12345 --lr 0.001 --batch-size 32 --epochs 200 --early-stopping 100 --riemannian_dist 

for the BCI-IV 2b dataset.

• \code\library: Riemannian embedding estimation, feature preprocessing and extraction files

• \code\model: Models for spatial, temporal and spatio-temporal streams of our architecture.

If you find this material useful, please cite the following article:

@article{zhang2020rfnet, 
  title={Spatio-Temporal EEG Representation Learning on Riemannian Manifold and Euclidean Space},
  author={Zhang, Guangyi and Etemad, Ali},
  journal={arXiv preprint arXiv:2008.08633},
  year={2020}
}

Should you have any questions, please feel free to contact me at [email protected].