This repository contains the official PyTorch implementation of the following paper:
DREAM: Domain Invariant and Contrastive Representation for Sleep Dynamics
abstract: Sleep staging is a key challenge in diagnosing and treating sleep-related diseases due to its labor-intensive, time-consuming, costly, and error-prone. With the availability of large-scale sleep signal data, many deep learning methods are proposed for automatic sleep staging. However, these existing methods face several challenges including the heterogeneity of patients' underlying health conditions; the infeasibility of extracting meaningful information from unlabeled sleep signal data to improve prediction performances; the difficulty modeling complex interactions between sleep stages; and the lack of an effective mechanism in clinical contexts that allows for human intervention when needed. In this paper, we propose a neural network architecture named DREAM to tackle these issues for automatic sleep staging. DREAM consists of (i) a feature representation network that generates robust representations for sleep signals via the variational auto-encoder framework and contrastive learning and (ii) a sleep stage classification network that explicitly models the interactions between sleep stages in the sequential context at both feature representation and label classification levels via Transformer and conditional random field architectures. Our experimental results indicate that DREAM significantly outperforms existing methods for automatic sleep staging on three sleep signal datasets. Further, DREAM provides an effective mechanism for quantifying uncertainty measures for its predictions, thereby helping sleep specialists intervene in cases of highly uncertain predictions, resulting in better diagnoses and treatments for patients in real clinical settings.
The proposed framework DREAM consists of two main networks: (i) Feature representation network that transforms the input sleep segment into sleep-relevant and subject-invariant feature representations based on the VAE framework and (ii) Sleep Stage Classification network that captures the dependencies between sleep segments in the sequential context to find the best corresponding sleep stage sequence.
In the first stage, the feature representation network is trained to capture the sleep-relevant and subject-invariant representation from one input sleep segment. In the second stage, the trained feature representation network with fixed weights is employed to generate the sequence of the representations for all segments of the input sequence, which is then used as input for the classification network to identify the corresponding sleep stage sequence.
Used modules: numpy, scipy, pandas, scikit-learn, TorchCRF, pytorch_metric_learning, and PyTorch (CUDA toolkit if use GPU).
$ conda create -n DREAM python=3.9.7
$ conda activate DREAM
$ conda install scipy pandas scikit-learn TorchCRF pytorch_metric_learning
$ conda install numpy=1.21.2
$ conda install pytorch torchvision torchaudio cudatoolkit=10.2 -c pytorch
Run the following scripts to download and preprocess SleepEDF-20 dataset
$ cd data/edf_20
$ bash download_edf_20.sh
$ cd ../../preprocess
$ python preprocess_edf.py
Run the following scripts to train DREAM
For supervised DREAM
$ batch job_batch.txt
For semi-supervised DREAM
$ batch job_batch_semi_sup.txt