Convolutional neural networks for EEG-based stress prediction.

This repository applies convolutional neural networks (CNNs) to identify stressed patients based on their electroencephalogram (EEG) recordings (i.e. brain signals). See writeup.pdf for a full description of the setup and results.

Required packages that aren't in the standard library are:

matplotlib, numpy, pandas, torch

This repository includes only the processed data (data/*) required for training, as the raw dataset is ~600MB. If you would like to process the raw data yourself using data_prep.ipynb, reach out to me for the full dataset (you'll need to set RAW_DATA_DIR appropriately).

train_test.ipynb loads the processed data, sets up CNN models and trains them to predict stress based on subintervals of EEG signals. The CNN architecture is defined in cnn.py and the batching algorithm in eeg_sampler.py. Helpful plotting functions are in perf_plots.py.

The ./results directory contains model parameters and training trajectories for the models from the writeup paper.

• .json files specifying model parameters
• .pt files storing the recorded training/test trajectory.

For references, see writeup.pdf.