SleepJEPA

Learning the latent world of sleep!

Overview

SleepJEPA is a Python package for at-home sleep study data to classify sleep stages, sleepiness, and estimate long-term disease risk. SleepJEPA was trained with the JEPA architecture, modified for signal data and sleep studies. The repository is built using nbdev, which means the package is developed in Jupyter notebooks.

See the publication in tbd…

Install

pip install sleepjepa

Inference

To perform inference on unseen data (using EDF file paths as input), use the sleepjepa_sleep_stage_inference.py and sleepjepa_sleep_stage_inference_config.yaml files.

Please create an account on Hugging Face and request access to the models here. You will also need to create a personal access token, with read access. Read more here.

Following, download the encoder and sleep stage classifier models with:

from sleepjepa.inference import download_sleepjepa_models

download_sleepjepa_models(models_dir='', token=YOUR_HF_TOKEN)

You will also be prompted to download the files when running the inference script, if they are not in models_dir.

After the models are downloaded, update the sleepjepa_sleep_stage_inference_config.yaml file with a path to your edf or edf directory. You can pass both a single edf file or a directory of edfs. You can also use glob syntax if specifying edfs within single sub directories (e.g. /path/to/base/directory/**/).

Unfortunately, due to differing naming conventions for signal channels, if you are passing multiple edfs, but they have different channel names, the dataloader will fail. It is recommended in this case to rewrite the edfs to a consistent channel name format, or perform inference on them one by one.

If a specific channel is not available for a given edf or set of edfs, pass the keyword “null” or “dummy” to that channels name parameter in the yaml if you’d like to see how the model performs with that channel set as all zero.

The model expects referenced EEG, Left EOG, EMG, and ECG channels. If your channels are unreferenced, you may pass the corresponding reference channels. The model was trained with the following referenced channels (priority was given to the ones listed first in the below list):

EEG: C4-M1 or C3-M2
EMG: Chin1-Chin2 or Chin1-Chin3
ECG: Augmented lead 2 (or ECG (LL) - ECG (RA))
Left EOG: E1-M2

Check the slumber.py source code for NSRR specific channels (under the SHHS_CHANNELS, MROS_CHANNELS, WSC_CHANNELS, etc. variables) if the above is confusing. Additionally bandpass, lowpass, and highpass filters need to be specified. The defaults are listed in the yaml and shouldn’t be changed.

For the device parameter, use “cpu” (slowest) or a GPU (e.g. “cuda:0”). MPS (“mps” for Mac OS X) is not supported due to the use of nested tensors.

Prediction logits of shape [bs x 5 x L] are outputted, where L is the total number of 30 second sleep epochs in the entire length sleep study (if your study is 8 hours long, then L will be 8 * 3600 divided by 30). The first dimension indicates individual sleep stage logits where the 0 index is wake, 1 is N1, 2 is N2, 3 is N3, and 4 is REM. To retrieve probabilities, use torch’s softmax function on the 1st dimension of the tensor. Note that the model uses nested tensors.

If using a CPU, it is recommended to predict each sample individually. This will happen automatically with the inference script by passing the dataset (not dataloader) to the infer_on_edf_dataset function. Ensure autocast is set to False when using a CPU.

If using CUDA, you can use a dataloader. On CUDA, autocast should be enabled and is an option in the inference function. If you are having issues, try inference on each item one by one by passing the dataset (not dataloader) to the infer_on_edf_dataset function with autocast set to False.

To finally run the predictions on a single edf file or directory of edf files, run the python script via:

python sleepjepa_sleep_stage_inference.py sleepjepa_sleep_stage_inference_config.yaml

Predictions will be output to a torch tensor file .pt at the location specified in the yaml.

Additionally, you can pass the save_hypjson parameter in the yaml as true. This will perform softmax, max index selecting, and save the predictions as a HYPJSON file (with the same filename as the edf file, + ’_sleepjepa.HYPJSON’). You can also use the write_pred_to_hypjson function for individual files. Sleep stages are mapped in this function to typical HYPJSON standards (for example, REM is mapped to the integer “5”).

Repository Structure and Usage

This is an nbdev repository, which means the package is developed in Jupyter notebooks located in the nbs/ directory. Any modifications or additions to the sleepjepa package should be made by editing these notebooks.

To build the package, run nbdev_prepare in the terminal. This will generate the sleepjepa package in the sleepjepa/ directory and all python modules, which can be imported and used in other Python projects.

To add new functionality, create a new notebook or add to exisitng in the nbs/ directory and follow the instructions in the nbdev documentation to add the new functionality. Then, run nbdev_prepare to generate the sleepjepa package with the new functionality.

Directory Structure:

• nbs/: Contains the source notebooks that generate the Python package
• jobs/: Contains processing and training scripts
  • apples/: Processing scripts for the apples dataset
  • human_sleep_project/: Processing scripts for the human sleep project dataset
  • jepa/: Training script for the JEPA sleep model
  • mesa/: Processing scripts for the mesa dataset
  • mnc/: Processing scripts for the mnc dataset
  • mros/: Processing scripts for the mros dataset
  • shhs/: Processing scripts for the shhs dataset
  • wsc/: Processing scripts for the wsc dataset
  • sleep_outcomes/:
    • config/: Config files for training outcomes/models
    • train_age.py: Model training script for SleepJEPA representations to estimate chronological age prediction
    • train_obj_sleepiness.py: Model training script for SleepJEPA representations to estimate objective sleepiness and narcolepsy outcomes
    • train_sleep_stages.py: Model training script for SleepJEPA representations to classify sleep stages
    • train_demographics_long_term_outcome.py: Model training script to train demographics only model for long term disease risk estimation
    • train_sleep_long_term_outcome.py: Model training script to train SleepJEPA representations +/- demographicss for long term disease risk estimation

Technical Details

• We trained the foundational model on H100nvl GPUs using PyTorch Lightning.
• We monitored training using the Weights and Biases platform.