This repository provides a PyTorch implementation of a submitted paper:
TransSleep: Transitioning-aware Attention-based Deep Neural Network for Automatic Sleep Staging
Jaeun Phyo1, Wonjun Ko1, Eunjin Jeon1, Heung-Il Suk1, 2
(1Department of Brain and Cognitive Engineering, Korea University)
(2Department of Artificial Intelligence, Korea University)Abstract: Sleep staging is essential for sleep assessment and plays a vital role as one of the health indicators. Due to its capability, recent studies exploited and devised various machine learning as well as deep learning architectures for sleep staging. However, two important challenges hinder their practical use. Specifically, effectively capturing salient waveforms in sleep signals and correctly classifying confusing stages in transitioning epochs are rarely easy. In this work, we propose a novel deep neural network structure that captures distinctive local temporal patterns and exploits two auxiliary tasks to distinguish confusing stages. Particularly, our proposed framework adopts an attention-based multi-scale feature extractor module to capture salient waveforms. A stage-confusion estimator module with a novel auxiliary task, epoch-level stage classification}, estimates confidence scores to identify confusing stages. A context encoder module with another novel stage-transition detection auxiliary task represents contextual relationships with transition information across neighboring epochs. With the estimated information by two auxiliary tasks as well as architectural strength, our proposed method shows promising results for the automatic sleep staging. We demonstrated the validity of our proposed method on two publicly available datasets Sleep-EDF and MASS, by achieving state-of-the-art performances. Moreover, we analyzed the results from a diverse perspective with ablative cases. Based on our results, we believe that the proposed network has immense potential to provide new insights into deep learning-based sleep staging.
To download Sleep-EDF dataset
To download MASS dataset
model.py contains the proposed deep learning architectures, utils.py contains functions used for experimental procedures, and main.py contains the main experimental functions.
This work was supported by Institute for Information & Communications Technology Promotion (IITP) grant funded by the Korea government under Grant 2017-0-00451 (Development of BCI based Brain and Cognitive Computing Technology for Recognizing User’s Intentions using Deep Learning) and Grant 2019-0-00079 (Department of Artificial Intelligence, Korea University).