Skip to content

CCRRSleepNet: A Hybrid Relational Inductive Biases Network for Automatic Sleep Stage Classification on Raw Single-Channel EEG

Notifications You must be signed in to change notification settings

nengwp/CCRRSleepNet

Repository files navigation

CCRRSleepNet

Code for the model in the paper CCRRSleepNet: A Hybrid Relational Inductive Biases Network for Automatic Sleep Stage Classification on Raw Single-Channel EEG by Neng Wenpeng, Lu Jun, and Xu Lei.

This work has been accepted for publication in Brain Sciences.

Prepare dataset

We evaluated our CCRRSleepNet with Sleep-EDF dataset.

For the Sleep-EDF dataset, you can run the following scripts to download SC subjects.

cd data
chmod +x download_physionet.sh
./download_physionet.sh

Then run the following script to extract specified EEG channels and their corresponding sleep stages.

python prepare_physionet.py --data_dir data --output_dir data/eeg_fpz_cz --select_ch 'EEG Fpz-Cz'
python prepare_physionet.py --data_dir data --output_dir data/eeg_pz_oz --select_ch 'EEG Pz-Oz'

Training a model

Run this script to train a DeepSleepNet model for the first fold of the 20-fold cross-validation.

python train.py --data_dir data/eeg_fpz_cz --output_dir output --n_folds 20 --fold_idx 0 --pretrain_epochs 100 --finetune_epochs 200 --resume False

You need to train a CCRRSleepNet model for every fold (i.e., fold_idx=0...19) before you can evaluate the performance. You can use the following script to run batch training

chmod +x batch_train.sh
./batch_train.sh data/eeg_fpz_cz/ output 20 0 19 0

Scoring sleep stages

Run this script to determine the sleep stages for the withheld subject for each cross-validation fold.

python predict.py --data_dir data/eeg_fpz_cz --model_dir output --output_dir output

The output will be stored in numpy files.

Get a summary

Run this script to show a summary of the performance of our CCRRSleepNet compared with the state-of-the-art hand-engineering approaches. The performance metrics are overall accuracy, per-class F1-score, and macro F1-score.

python summary.py --data_dir output

Acknowledgement

We refer to part of the code of DeepSleepNet

About

CCRRSleepNet: A Hybrid Relational Inductive Biases Network for Automatic Sleep Stage Classification on Raw Single-Channel EEG

Topics

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published