Automatic sleep staging by a hybrid model based on deep 1D-ResNet-SE and LSTM with single-channel raw EEG signals

The SE module

The architecture of the SE module is depicted in Fig. 1. The SE module is designed to calibrate channel-wise features by considering inter-channel dependencies, consisting of three operations, namely squeezing, excitation, and scaling. Specifically, multiple feature maps produced by the convolutional operation serve as input to the SE module. In the squeezing operation, the feature maps are squeezed by the global average pooling (GAP) layer (Szegedy et al., 2015) to create a weight vector. Next, in the excitation operation, the weight vector is passed through two successive fully connected (FC) layers, each of which is activated by a rectified linear unit (ReLU) (Nair & Hinton, 2010) and a sigmoid function, respectively. These procedures introduce a nonlinear transformation and provide a smooth gating calculation, resulting in a channel weight vector. Finally, the channel weight vector is multiplied by the input feature maps to produce the final output in the scaling operation. In other words, the input feature maps are reweighted by the channel weight vector.

The 1D-ResNet-SE block

The proposed 1D-ResNet-SE block, as illustrated in Fig. 2, consists of two successive one-dimensional convolutional (Conv1D) layers. The first Conv1D layer is followed by a batch normalization (BN) layer and a ReLU activation function layer, while the second Conv1D layer is followed by a BN layer, an SE module, and a ReLU activation function layer sequentially. The Conv1D layer is utilized to extract relevant features from the raw data (Kiranyaz et al., 2021), while the BN layer is employed to standardize the data batch and improve the generalization capability of the model, also addressing the common issue of vanishing gradient in deep learning (Ioffe & Szegedy, 2015). The ReLU activation function is a nonlinear function that enhances the model’s expression ability. The SE module is incorporated after the second BN layer to calibrate the features.

To tackle the challenges of gradient dispersion and explosion, a shortcut connection is established between the input x and output of the SE module. The shortcut path involves a Conv1D layer with a convolutional kernel size of one to match the input and output dimensions, followed by a BN layer. For a network with input x, the underlying mapping is denoted as H(x), and the residual mapping is $F\left(x\right)=H\left(x\right)-x$. The rationale behind the residual network is to enable the neural network to fit the residual mapping instead of learning the underlying mapping. The output y from the shortcut operation is computed as $y=F\left(x\right)+x$.

LSTM network

The LSTM network, a variant of RNNs, was proposed by Hochreiter & Schmidhuber (1997). This network is specifically designed to capture long-term dependencies and address the issues of vanishing and exploding gradients commonly encountered during long-sequence learning. In contrast to ordinary RNNs, the LSTM network includes the memory cell instead of the recurrent node.

At each time step, the LSTM network inputs the current time step input X_t and the hidden state of the previous time step H_t−1. The network comprises three gates, namely forget, input, and output gates, responsible for information selection. These gates produce output values, F_t, I_t, and O_t, respectively, using a fully connected layer with the sigmoid σ activation function, ranging between 0 and 1. Among these gates, the forget gate decides which information from the previous memory cell’s internal state C_t−1 needs to be removed or retained. The input gate determines which information in the input node $\widetilde{C_t}$ should be added to the current memory cell’s internal state C_t, where the input node is obtained from a fully connected layer with the tanh activation function that ranges between −1 and 1. The output gate decides which information in the updated cell state should be included in the output hidden state at the current time step H_t. Figure 3 illustrates the LSTM network and its internal architecture.

The proposed model

The proposed model, 1D-ResNet-SE-LSTM, aims to predict a sequence of sleep stages in a many-to-many scheme from a series of EEG epochs. The input to the model consists of an ordered set of N single-channel EEG epochs, denoted by $\left\{x_1,\dots ,x_N\right\}$, where x_i ∈ ℝ^E_s×F_s, x_i represents the i-th EEG epoch, and 1 ≤ i ≤ N, with E_s representing the duration of the EEG epoch in seconds, and F_s denoting the EEG signals’ sampling rate. The model generates an equally long sequence of predicted sleep stages $\left\{{\hat{y}}_1,\dots ,{\hat{y}}_N\right\}$, where ${\hat{y}}_i$ corresponds to the predicted sleep stage for x_i, and ${\hat{y}}_i\in \left\{0,1,2,3,4\right\}$ denotes the five sleep stages W, N1, N2, N3, and REM, respectively. The value of N is set to 15, E_s is 30, and F_s is 100 for this study.

The model’s architecture consists of three sequential 1D-ResNet-SE blocks, each with 32, 64, and 128 convolutional kernels, and convolutional kernel sizes of 3, 5, and 7, respectively. The strides are set to 1, and no padding is applied. The first two blocks are followed by a max-pooling layer that reduces the feature map size by a factor of four and preserves essential features, while the last block is followed by a GAP layer that averages each feature map. A dropout layer with a regularization probability of 50% is applied to improve the model’s generalization ability (Srivastava et al., 2014). The 1D-ResNet-SE blocks serve as the epoch-wise feature extractor to convert the EEG epoch x_i into the feature vector a_i using learnable parameters θ_r, as given by Eq. (1). (1)${\displaystyle a_i=ResNe{t}_{{\theta }_r}\left(x_i\right)}$

As a result, an input sequence of EEG epochs is transformed into a sequence of feature vectors, which is then fed into the LSTM layer as a series of inputs, where each feature vector corresponds to a specific time step in the sequence. The LSTM layer has 64 hidden units and is responsible for learning transition rules among sleep stages. Given N feature vectors $\left\{a_1,\dots ,a_N\right\}$ arranged sequentially, the LSTM layer processes the i-th feature vector a_i using Eq. (2): (2)${\displaystyle h_i,c_i=LST{M}_{{\theta }_s}\left(h_{i-1},c_{i-1},a_i\right)}$

where LSTM_θs represents the LSTM layer processing the sequence of feature vectors a_i, θ_s is the learnable parameters of the LSTM, h_i and c_i are vectors of hidden and cell states of the LSTM layer after processing the feature vector a_i, h_i−1, and c_i−1 are the hidden and cell states from the previous time step, h₀ and c₀ are the initial hidden and cell states that are set to 0.

Lastly, a Softmax layer is employed for the five-class classification, producing a five-dimensional output vector corresponding to the probabilities of each class (W, N1, N2, N3, and REM) that the given EEG epoch belongs to. The sleep stage label for each epoch is determined by the maximum class probability using the argmax function, resulting in the model generating a predictive sleep stage sequence. The model’s schematic representation is shown in Fig. 4.

Loss function

To address the issue of imbalanced class distribution in the dataset, the weighted cross-entropy loss function was utilized in this study. The weight coefficients are determined based on the number of samples in each class and the difficulty level of distinguishing each class. The N1 and N3 stages, which have fewer samples than the other stages, are considered minority classes. Moreover, the N1 stage is particularly challenging to classify due to its transitional nature from wakefulness to sleeping (Phan et al., 2018), whereas the N3 stage has distinctive characteristics, such as significantly higher amplitude, making it easier to identify. To account for these imbalances, a weight coefficient of 1.5 was assigned to the N1 stage, while the weight coefficients for the remaining four stages were set to 1. The weighted cross-entropy loss function is formulated as Eq. (3). (3)${\displaystyle Loss=-\sum _i^C{w}_i\ast y_i\ast log\left({\hat{y}}_i\right)}$

where y_i and ${\hat{y}}_i$ represent the actual and predicted labels, respectively, for each class i in C classes, and w_i denotes the weight coefficient assigned to each class.

Training setup

The dataset was partitioned into training and test sets with proportions of 90% and 10%, respectively, while 10% of the training set was extracted as the validation set. To minimize the loss function, the adaptive moment estimation (Adam) optimizer (Kingma & Ba, 2014) was employed, with an initial learning rate of 0.001. The step size was reduced by a factor of 10 after every 10 epochs if the validation loss did not decrease. To prevent overfitting, early stopping was implemented when the validation loss remained unchanged for 20 epochs (Basheer & Hajmeer, 2000). The model with the highest validation accuracy was selected as the final model based on the model checkpoints. A training batch size of 16 was used, and the model was trained for 100 epochs. Finally, the performance of the model was assessed using the test set.

Model implement environment

We implemented our model using the TensorFlow and Keras frameworks in a Python 3 environment running on an Ubuntu 20.04 system. The computation was performed on an NVIDIA GRID T4-4Q GPU with 4GB of GPU memory, hosted on the Alibaba Cloud Elastic Compute Service.

Per class evaluation

The model’ performance was evaluated by calculating precision, recall, and F1-score for each class on the test set. These metrics are defined as follows:

(4)${\displaystyle Precision=\frac{TP}{TP+FP},}$ (5)${\displaystyle Recall=\frac{TP}{TP+FN},}$ (6)${\displaystyle F1-score=\frac{2\times TP}{2\times TP+FP+FN},}$

where TP, FP, TN, and FN refer to the number of true positive, false positive, true negative, and false negative predictions, respectively.

Overall evaluation

The model’s overall performance was evaluated using the overall accuracy (Acc), as defined in Eq. (7), where C = 5 represents the number of target classes, TP_i denotes the number of true positives in class i, and M is the total number of EEG epochs. (7)${\displaystyle Acc=\frac{\sum _{i=1}^{C}T{P}_i}{M}}$

To account for the imbalanced nature of the datasets, the macro average F1-score (MF1) was also used, which is defined in Eq. (8), where F1_i is per-class F1-score for class i. (8)${\displaystyle MF1=\frac{\sum _{i=1}^{C}F{1}_i}{C}}$

Moreover, Cohen’s kappa coefficient (κ) (Cohen, 1960) was computed to assess the degree of agreement between the model’s classification results and those of human sleep experts, as given by Eq. (9): (9)${\displaystyle \kappa =\frac{p_o-p_e}{1-p_e}}$

where p_o denotes the number of samples the raters agree on and is divided by the total number of samples, p_e represents the probability of chance agreement.

In addition, the confusion matrix and its normalized version were generated to provide insight into the classification performance.