A 1D CNN for high accuracy classification and transfer learning in motor imagery EEG-based brain-computer interface
In order to train the network (and perform inference) you need to install the dependencies. There are two ways to install dependencies: (1) Using a docker container (recommended), (2) Using a python environment. An NVIDIA GPU with at least 6 GB is also recommended. The network was trained with an NVIDIA RTX 2060 and an NVIDIA TESLA P100, the former taking about 30 minutes to train, the latter about 15 minutes.
Assuming you have docker installed. Building the container is straightforward.
- Clone this repository
git clone https://github.com/Kubasinska/MI-EEG-1D-CNN.git
- Open a shell and
cd
into the docker folderMI-EEG-1D-CNN/docker
- Run
docker build -t eegcnn .
Depending on the permissions set on your machine you may need to run this command as root by addingsudo
at the beginning. This command will build a new container calledeegcnn
. This may take a few minutes. During the build process, the original dataset is downloaded and generated as described in the methods section of the original paper. Once the container is built, you already have everything ready. - Run the container to check that everything has been installed correctly. To launch the container, go to the working directory of the
MI-EEG-1D-CNN
repository. You can launch the container in several ways; here, we recommend 2, one that allows you to plot graphs and one that does not.- I want to see some graphs. If you want to see graphs, you must allow the container to access your screen To do this, run
xhost +local:root
; this exposes your xhost so that the container can render the correct display by reading and writing through the X11 unix socket. Then, from the working directory of the repository launch:docker run -it --gpus all -v $(pwd):/workspace -v /tmp/.X11-unix:/tmp/.X11-unix:rw -e DISPLAY=unix$DISPLAY --device /dev/dri --privileged -v /home/$USER/.Xauthority:/root/.Xauthority eegcnn bash
. This will open a bash shell inside the container, and you are good to go! If you don't have an NVIDIA GPU or have NVIDIA Container Runtime (v2) installed, omit--gpus all
. When you close the container, remember to launchxhost -local:root
. - I don't care about graphs. In this case, it is much easier! Just run:
docker run -it --gpus all -v $(pwd):/workspace eegcnn bash
and you're in.
- I want to see some graphs. If you want to see graphs, you must allow the container to access your screen To do this, run
NOTE: All code paths are made, so you don't have to change anything using the container. For example, inside the directory
/dataset
of the container, you find three sub-folders/dataset/original
, the original dataset data,/dataset/paper
the data generated through the method described in the paper with the scriptdataset_generator/generator.py
, and a third folder/dataset/saved_models
, where the trained models automatically save after training. When you run the script to make an inference, you don't have to change anything because the python interpreter already knows where to get the trained model.
This procedure is more straightforward but can create dependency issues based on your machine or operating system. Using this procedure, you must manually download the original dataset and generate the dataset used in the paper. You also need Anaconda or Miniconda to create a separate python environment. The following guide assumes you have Anaconda or Miniconda installed on your system.
- Open a terminal and
cd
into theMI-EEG-1D-CNN/docs
folder. Runconda env create -f environment.yml
. These create a new python environment calledeeg
containing almost all necessary dependencies. The only dependencies missing are CUDA and cuDNN, which TensorFlow needs to use your GPU. If you don't have an NVIDIA GPU, go ahead. If you have an NVIDIA GPU, you need to install CUDA 10.1 and cuDNN 7.6 (be careful with the version, CUDA 10.2 or 10.0 is not good, you need 10.1, same for cuDNN). Please refer to the official NVIDIA website for installation, here is a guide for windows. - Download the EEG Motor Movement/Imagery Dataset here. The dataset is quite large (3.4 GB); it will take a while. Once downloaded, extract it. If you have wget you can download it from the terminal with the command.
wget -r -N -c -np https://physionet.org/files/eegmmidb/1.0.0/
. - Generate the dataset; this procedure simply takes the raw data and breaks it into the input dimension of the neural network. Use the script
MI-EEG-1D-CNN/dataset_generator/generator.py
. Change the dataset path to the path of the dataset you downloaded, and you are ready! Don't forget to run the script with the new conda environmenteeg
.
From here on, we assume that all dependencies have been installed correctly.
.
└── MI-EEG-1D-CNN/
├── data_processing/ # A module with useful functions
│ └── general_processor.py
├── dataset_generator/ # Script that generates the dataset
│ └── generator.py
├── docker/ # All the useful things to build the container
│ ├── Dockerfile
│ ├── environment.yml
│ └── generator.py
├── docs/ # Scripts for inference and plotting
│ ├── inference
│ └── environment.yml
├── models/ # Scripts for training networks
│ ├── hand_test
│ └── transfer
└── model_set/ # A module with all models
└── models.py
The data_processing
folder is a module that contains many helpful functions for assigning labels, generating the dataset as described in the paper, and other essential things. The dataset_generator
folder contains the script that generates the dataset. Running the script generator.py
with the correct paths creates a new dataset described in the paper. The script saves the data of each subject and channel combinations separately, so it is possible to load the data subject by subject into memory. The docker
folder contains the Dockerfile and the environment.yml
file. The docs
folder contains the scripts for inference and plotting. Inside the inference subfolder, you can find one script for each ROI described in the paper (including the test without batch normalization and the test without SMOTE). There are also scripts to do transfer learning on the seven random subjects. Be careful; these scripts assume that there is a trained model. If you want to run these scripts, you must first train the models and save them. The scripts to train the models can be found in the folder models
. Inside this folder, you will find a script for each ROI. The name pattern is train_
and the corresponding letter of each ROI (a,b,c,d,e,f
). For example, running the script train_a.py
will start training the network with the ROI-A, composed of channels ( [FC1, FC2], [FC3, FC4], [FC5, FC6]).
Be careful; if you are using the container, you won't have any problems; the script already knows where to get the data and save the model. If you are not using the docker container, you will have to change the paths manually.
source_path
refers to the dataset generated through thegenerator.py
script.save_path
refers to the path where the trained model will be saved.
The hand_test
subfolder contains the tests done on the network to evaluate the importance of data augmentation, checkpointing, and batch normalization. As usual, if you don't use our container, change the paths. The transfer
subdirectory contains seven scripts that train the model with ROI-E and exclude a single subject. These scripts are used to evaluate transfer learning. The model_set
folder contains the neural network (informally called HopefulNet) written in tensorflow using the subclass api.
@article{mattioli20211d,
title={A 1D CNN for high accuracy classification and transfer learning in motor imagery EEG-based brain-computer interface},
author={Mattioli, Francesco and Porcaro, Camillo and Baldassarre, Gianluca},
journal={Journal of Neural Engineering},
year={2021},
publisher={IOP Publishing}
}