caueeg-ceednet

CEEDNet: CAUEEG End-to-end Deep neural Network for automatic early detection of dementia based on the CAUEEG dataset.

This repository is the official implementation of "Deep learning-based EEG analysis to classify mild cognitive impairment for early detection of dementia: algorithms and benchmarks" from the CNIR (CAU NeuroImaging Research) team.

Model summary

CAUEEG-Dementia dataset

CAUEEG-Abnormal dataset

Getting started

Requirements

• Installation of Conda (refer to https://www.anaconda.com/products/distribution)
• Nvidia GPU with CUDA support

Note: we tested the code in the following environments.

OS	Python	PyTorch	CUDA
Windows 10	3.9.12	1.11.0	11.3
Ubuntu 18.04	3.8.11	1.10.0	11.3
Ubuntu 20.04	3.9.12	1.11.0	11.3

Installation

(optional) Create and activate a Conda environment.

conda create -n caueeg python=3.9
conda activate caueeg

Install PyTorch library (refer to https://pytorch.org/get-started/locally/).

conda install pytorch torchvision cudatoolkit=11.3 -c pytorch

Install other necessary libraries.

pip install -r requirements.txt

Preparation of the CAUEEG dataset

❗ Note: The use of the CAUEEG dataset is allowed for only academic and research purposes 👩‍🎓👨🏼‍🎓.

• For full access of the CAUEEG dataset, follow the instructions specified in https://github.com/ipis-mjkim/caueeg-dataset.
• Download, unzip, and move the whole dataset files into local/datasets/.

💡 Note: We provide caueeg-dataset-test-only at [link 1] or [link 2] to test our research. caueeg-dataset-test-only has the 'real' test splits of two benchmarks (CAUEEG-Dementia and CAUEEG-Abnormal) but includes the 'fake' train and validation splits.

Usage

Train

Train a CEEDNet model on the training set of CAUEEG-Dementia from scratch using the following command:

python run_train.py data=caueeg-dementia model=1D-ResNet-18 train=base_train

Similarly, train a model on the training set of CAUEEG-Abnormal from scratch using:

python run_train.py data=caueeg-abnormal model=1D-ResNet-18 train=base_train

Or, you can use this Jupyter notebook.

If you encounter a GPU memory allocation error or wish to adjust the balance between memory usage and training speed, you can specify the minibatch size by adding the ++model.minibatch=INTEGER_NUMBER option to the command as shown below:

python run_train.py data=caueeg-dementia model=1D-ResNet-18 train=base_train ++model.minibatch=32

python run_train.py data=caueeg-abnormal model=1D-ResNet-18 train=base_train ++model.minibatch=32

Thanks to Hydra support, the model, hyperparameters, and other training details are easily tuned using or modifying config files.

python run_train.py data=caueeg-dementia model=2D-VGG-19 train=base_train

For speed-up, we recommend using the PyArrow.feather file format than using directly EDF, which can be done:

python ./datasets/convert_file_format.py  # it takes a few minutes
python run_train.py data=caueeg-dementia model=2D-VGG-19 train=base_train ++data.file_format=feather

Evaluation

Evaluation can be conducted using this Jupyter notebook (or another notebook for caueeg-dataset-test-only case)

To use the pre-trained model, download the checkpoint file from here, and move it to local/checkpoint directory (e.g., local/checkpoint/1vc80n1f/checkpoint.pt for the 1D-VGG-19 model on the CAUEEG-Dementia benchmark).

Citation

If you found this repository helpful, please cite the paper below.

@article{kim2023deep,
  title={Deep learning-based EEG analysis to classify normal, mild cognitive impairment, and dementia: algorithms and dataset},
  author={Kim, Min-jae and Youn, Young Chul and Paik, Joonki},
  journal={NeuroImage},
  pages={120054},
  year={2023},
  publisher={Elsevier}
}