SleepEGAN: A GAN-enhanced ensemble deep learning model for imbalanced classification of sleep stages

Abstract

Deep neural networks have played an important role in the automatic classification of sleep stages due to their strong representation and in-model feature transformation abilities. However, class imbalance and individual heterogeneity which typically exist in raw EEG signals of sleep data can significantly affect the classification performance of any machine learning algorithms. To solve these two problems, this paper develops a generative adversarial network (GAN)-powered ensemble deep learning model, named SleepEGAN, for the imbalanced classification of sleep stages. To alleviate class imbalance, we propose a new GAN (called EGAN) architecture adapted to the features of EEG signals for data augmentation. The generated samples for minority classes are used in the training process. In addition, we design a cost-free ensemble learning strategy to reduce the model estimation variance caused by the heterogeneity between the validation and test sets, to enhance the accuracy and robustness of prediction performance. We show that the proposed method improves classification accuracy compared to several existing state-of-the-art methods. The overall classification accuracy and macro F1-score obtained by our SleepEGAN method on three public sleep datasets are: Sleep-EDF-39: 86.8% and 81.9%; Sleep-EDF-153: 83.8% and 78.7%; SHHS: 88.0% and 82.1%.