Initial commit

Author: yangdai97

.gitattributes

@@ -0,0 +1,2 @@
+# Auto detect text files and perform LF normalization
+* text=auto

Kfold_trainer.py

@@ -0,0 +1,152 @@
+import os
+import numpy as np
+from tqdm import tqdm
+
+import torch
+from torch import nn
+from torch import optim
+from torch.autograd import Variable
+from torch.utils.data import TensorDataset, DataLoader
+
+from sklearn.model_selection import StratifiedKFold
+from sklearn.metrics import accuracy_score
+
+from model import Transformer
+from early_stop_tool import EarlyStopping
+from data_loader import data_generator
+from args import Config, Path
+
+
+def set_random_seed(seed=0):
+    np.random.seed(seed)
+    torch.manual_seed(seed)        # CPU
+    torch.cuda.manual_seed(seed)   # GPU
+
+
+def test(model, test_loader, config):
+    criterion = nn.CrossEntropyLoss()
+    model.eval()
+
+    pred = []
+    label = []
+
+    test_loss = 0
+
+    with torch.no_grad():
+        for batch_idx, (data, target) in enumerate(test_loader):
+            data = data.to(config.device)
+            target = target.to(config.device)
+            data, target = Variable(data), Variable(target)
+
+            output = model(data)
+            test_loss += criterion(output, target.long()).item()
+
+            pred.extend(np.argmax(output.data.cpu().numpy(), axis=1))
+            label.extend(target.data.cpu().numpy())
+
+        accuracy = accuracy_score(label, pred, normalize=True, sample_weight=None)
+
+    return accuracy, test_loss
+
+
+def train(save_all_checkpoint=False):
+    config = Config()
+    path = Path()
+
+    dataset, labels, val_loader = data_generator(path_labels=path.path_labels, path_dataset=path.path_TF)
+
+    kf = StratifiedKFold(n_splits=config.num_fold, shuffle=True, random_state=0)
+
+    for fold, (train_idx, test_idx) in enumerate(kf.split(dataset, labels)):
+        print('\n', '-' * 15, '>', f'Fold {fold}', '<', '-' * 15)
+        if not os.path.exists('./Kfold_models/fold{}'.format(fold)):
+            os.makedirs('./Kfold_models/fold{}'.format(fold))
+
+        X_train, X_test = dataset[train_idx], dataset[test_idx]
+        y_train, y_test = labels[train_idx], labels[test_idx]
+        train_set = TensorDataset(X_train, y_train)
+        test_set = TensorDataset(X_test, y_test)
+        train_loader = DataLoader(dataset=train_set, batch_size=config.batch_size, shuffle=False)
+        test_loader = DataLoader(dataset=test_set, batch_size=config.batch_size, shuffle=False)
+
+        model = Transformer(config)
+        model = model.to(config.device)
+
+        criterion = nn.CrossEntropyLoss()
+
+        # AdamW optimizer
+        optimizer = optim.AdamW(model.parameters(), lr=config.learning_rate, weight_decay=0.01)
+
+        # apply early_stop. If you want to view the full training process, set the save_all_checkpoint True
+        early_stopping = EarlyStopping(patience=20, verbose=True, save_all_checkpoint=save_all_checkpoint)
+
+        # evaluating indicator
+        train_ACC = []
+        train_LOSS = []
+        test_ACC = []
+        test_LOSS = []
+        val_ACC = []
+        val_LOSS = []
+
+        for epoch in range(config.num_epochs):
+            running_loss = 0.0
+            correct = 0
+
+            model.train()
+
+            loop = tqdm(enumerate(train_loader), total=len(train_loader))
+            for batch_idx, (data, target) in loop:
+                data = data.to(config.device)
+                target = target.to(config.device)
+                data, target = Variable(data), Variable(target)
+
+                optimizer.zero_grad()
+                output = model(data)
+
+                loss = criterion(output, target.long())
+
+                loss.backward()
+
+                optimizer.step()
+
+                running_loss += loss.item()
+
+                train_acc_batch = np.sum(np.argmax(np.array(output.data.cpu()), axis=1) == np.array(target.data.cpu())) / (target.shape[0])
+                loop.set_postfix(train_acc=train_acc_batch, loss=loss.item())
+                correct += np.sum(np.argmax(np.array(output.data.cpu()), axis=1) == np.array(target.data.cpu()))
+
+            train_acc = correct / len(train_loader.dataset)
+            test_acc, test_loss = test(model, test_loader, config)
+            val_acc, val_loss = test(model, val_loader, config)
+            print('Epoch: ', epoch,
+                  '| train loss: %.4f' % running_loss, '| train acc: %.4f' % train_acc,
+                  '| val acc: %.4f' % val_acc, '| val loss: %.4f' % val_loss,
+                  '| test acc: %.4f' % test_acc, '| test loss: %.4f' % test_loss)
+
+            train_ACC.append(train_acc)
+            train_LOSS.append(running_loss)
+            test_ACC.append(test_acc)
+            test_LOSS.append(test_loss)
+            val_ACC.append(val_acc)
+            val_LOSS.append(val_loss)
+
+            # Check whether to continue training. If save_all_checkpoint=False, the model name will be ‘model.pkl'
+            early_stopping(val_acc, model, path='./Kfold_models/fold{}/model_{}_epoch{}.pkl'.format(fold, fold, epoch))
+
+            if early_stopping.early_stop:
+                print("Early stopping at epoch ", epoch)
+                break
+
+        np.save('./Kfold_models/fold{}/train_LOSS.npy'.format(fold), np.array(train_LOSS))
+        np.save('./Kfold_models/fold{}/train_ACC.npy'.format(fold), np.array(train_ACC))
+        np.save('./Kfold_models/fold{}/test_LOSS.npy'.format(fold), np.array(test_LOSS))
+        np.save('./Kfold_models/fold{}/test_ACC.npy'.format(fold), np.array(test_ACC))
+        np.save('./Kfold_models/fold{}/val_LOSS.npy'.format(fold), np.array(val_LOSS))
+        np.save('./Kfold_models/fold{}/val_ACC.npy'.format(fold), np.array(val_ACC))
+
+        del model
+
+
+if __name__ == '__main__':
+    set_random_seed(0)
+    train(save_all_checkpoint=False)

README.md

@@ -0,0 +1,52 @@
+# MultiChannelSleepNet
+### MultiChannelSleepNet: A Transformer-Based Deep Learning Method for Automatic Sleep Stage Classification With Multi-Channel PSG 
+#### *by: Yang Dai, Shanshan Liang, Lukang Wang, Qingtian Duan, Hui Yang, Chunqing Zhang\*, Xiaowei Chen\*, and Xiang Liao\* 
+
+
+## Abstract
+![AttnSleep Architecture](imgs/MutiChannelSleepNet.png)
+Automatic sleep staging aims to classify sleep states by machine intelligence, which plays an essential role in the measurement of sleep quality and diagnosis of sleep disorders. Although many automatic methods have been proposed for accurate sleep staging, most of these approaches utilize only single channel EEG signals to classify sleep stages. Multi-channel polysomnography (PSG) data provide more information to train classifier and may achieve higher sleep staging performance. Here we propose a novel transformer encoder-based deep learning method named MultiChannelSleepNet for sleep stage classification with multi-channel PSG data. The proposed model architecture is implemented based on the transformer encoder for single-channel feature extraction and multi-channel feature fusion. In each single-channel feature extraction block, transformer encoders are used to extract features from time-frequency maps of each channel. Based on our integration strategy, the feature maps extracted from each channel are fused for further feature extraction. Transformer encoders and a residual connection are applied in this module to preserve original information of each channel. The experimental results on three public datasets show that our method outperforms state-of-the-art methods in terms of different evaluation metrics.
+
+
+## Requirmenets:
+- python3.6
+- pytorch=='1.9.1'
+- numpy
+- sklearn
+- scipy=='1.5.4'
+- mne=='0.23.4'
+- tqdm
+
+## Data
+We used three public datasets in this study:
+
+- SleepEDF-20 (The first 39 records in SleepEDF-78)
+- [SleepEDF-78](https://physionet.org/content/sleep-edfx/1.0.0/)
+- [SHHS](https://sleepdata.org/datasets/shhs)
+
+This project currently only provides pre-processing code for SleepEDF-20 and SleepEDF-78. We will update the code of SHHS later.  
+After downloading the datasets, please place them in the folder with the corresponding name in the directory `dataset`.  
+You can run the `dataset_prepare.py` to extract events from the original record (.edf)
+
+## Reproducibility
+If you want to update the training parameters, you can edit the `args.py` file. In this file, you can update:
+
+- Device (GPU or CPU).
+- Batch size.
+- Number of folds (as we use K-fold cross validation).
+- The number of training epochs.
+- Parameters in our model (dropout rate, number of transformer encoder, etc)
+
+To easily reproduce the results you can follow the next steps:  
+
+1. Run `dataset_prepare.py` to extract events from the original record (.edf).
+2. Run `data_preprocess_TF` to preprocess the data. The original signals will be converted to time-frequency images, and normalized.
+3. Run `Kfold_trainer.py` to perform the standard K-fold cross validation.
+4. Run `result_evaluate.py` to get the evaluation report. It concludes the various valuation metrics we described in paper.  
+
+
+## Contact
+Yang Dai  
+Center for Neurointelligence, School of Medicine  
+Chongqing University, Chongqing 400030, China  
+Email: valar_d@163.com

args.py

@@ -0,0 +1,41 @@
+import torch
+import os
+
+
+class Config(object):
+    """args in model and trainer"""
+    def __init__(self):
+        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+        self.num_fold = 10
+        self.num_classes = 5
+        self.num_epochs = 200               # Because early stopping is used, this parameter can be relatively large
+        self.batch_size = 64
+        self.pad_size = 29                  # time dimension of TF image
+        self.learning_rate = 5e-6
+        self.dropout = 0.1                  # dropout rate in transformer encoder
+        self.dim_model = 128                # frequency of TF image
+        self.forward_hidden = 1024          # hidden units of transformer encoder
+        self.fc_hidden = 1024               # hidden units of FC layers
+        self.num_head = 8
+        self.num_encoder = 16               # number of encoders in single-channel feature extraction block
+        self.num_encoder_multi = 4          # number of encoders in multi-channel feature fusion block
+
+
+class Path(object):
+    """path of files in this project"""
+    def __init__(self):
+        self.path_PSG = 'dataset/sleepEDF-78/sleep-cassette'
+        self.path_hypnogram = 'dataset/sleepEDF-78/Hypnogram'
+        self.path_raw_data = 'data/sleepEDF-78/data_array/raw_data'
+        self.path_labels = 'data/sleepEDF-78/data_array/raw_data/labels'
+        self.path_TF = 'data/sleepEDF-78/data_array/TF_data'
+
+        if not os.path.exists(self.path_hypnogram):
+            os.makedirs(self.path_hypnogram)
+
+        if not os.path.exists(self.path_raw_data):
+            os.makedirs(self.path_raw_data)
+
+        if not os.path.exists(self.path_TF):
+            os.makedirs(self.path_TF)
+

data_loader.py

@@ -0,0 +1,41 @@
+import os
+import numpy as np
+
+import torch
+from torch.utils.data import TensorDataset, DataLoader
+from sklearn.model_selection import train_test_split
+
+from args import Config, Path
+
+
+def data_generator(path_labels, path_dataset):
+    config = Config()
+    dir_annotation = os.listdir(path_labels)
+
+    first = True
+    for f in dir_annotation:
+        if first:
+            labels = np.load(os.path.join(path_labels, f))
+            first = False
+        else:
+            temp = np.load(os.path.join(path_labels, f))
+            labels = np.append(labels, temp, axis=0)
+    labels = torch.from_numpy(labels)
+
+    dataset_EEG_FpzCz = np.load(os.path.join(path_dataset, 'TF_EEG_Fpz-Cz_mean_std.npy')).astype('float32')
+    dataset_EEG_PzOz = np.load(os.path.join(path_dataset, 'TF_EEG_Pz-Oz_mean_std.npy')).astype('float32')
+    dataset_EOG = np.load(os.path.join(path_dataset, 'TF_EOG_mean_std.npy')).astype('float32')
+
+    dataset = np.stack((dataset_EEG_FpzCz, dataset_EEG_PzOz, dataset_EOG), axis=1)
+    dataset = torch.from_numpy(dataset)
+
+    print('dataset: ', dataset.shape)
+
+    # hold out the validation set
+    X_train_test, X_val, y_train_test, y_val = train_test_split(dataset, labels, test_size=1/(config.num_fold+1), random_state=0, stratify=labels)
+
+    val_set = TensorDataset(X_val, y_val)
+    val_loader = DataLoader(dataset=val_set, batch_size=config.batch_size, shuffle=False)
+
+    print('val_set:', len(X_val))
+    return X_train_test, y_train_test, val_loader

data_preprocess_TF.py

@@ -0,0 +1,97 @@
+import os
+import numpy as np
+
+from scipy.fftpack import fft
+from scipy import signal
+from tqdm import tqdm
+
+from args import Path
+
+
+def data_array_concat(path_array):
+    """concat data from each subject"""
+    dir_PSG = os.listdir(path_array)
+    first = True
+    print('Preparing dataset:')
+    for f in tqdm(dir_PSG):
+        if first:
+            data_channel = np.load(os.path.join(path_array, f)).astype('float32')
+            first = False
+        else:
+            temp = np.load(os.path.join(path_array, f)).astype('float32')
+            data_channel = np.append(data_channel, temp, axis=0)
+    data_channel = np.squeeze(data_channel, axis=1)
+    return data_channel
+
+
+def spectrogram(x, window, n_overlap, nfft):
+    """
+    Transform to time-frequency images. This function imitates function spectrogram in Matlab
+    Args:
+        x (numpy array): Data
+        window (int): Size of window function
+        n_overlap (int):Number of coincidence points between two segments
+        nfft (int): Number of points during Fast Fourier Transform
+    """
+    len_x = len(x)
+    step = window - n_overlap
+    nn = nfft // 2 + 1
+    num_win = int(np.floor((len_x - n_overlap) / (window - n_overlap)))
+    spectrogram_data = []
+    # Hamming window default
+    win = signal.hamming(window)
+    for i in range(num_win):
+        subdata = x[i * step: i * step + window]
+        F = fft(subdata * win, n=nfft)
+        spectrogram_data.append(F[:nn])
+    spectrogram_data = np.array(spectrogram_data)
+    return spectrogram_data
+
+
+def data_normalize(dataset, channel):
+    """normalize datasets of each channel to zero mean and unit variance"""
+    for i in tqdm(range(dataset.shape[0])):
+        if True in np.isinf(dataset[i]):
+            for j in range(29):
+                if True in np.isinf(dataset[i][j]):
+                    for k in range(128):
+                        if np.isinf(dataset[i][j][k]):
+                            if k != 127:
+                                print('location of inf: ', i, ',', j, ',', k)
+                            if k == 0:
+                                if j == 0:
+                                    dataset[i][j][k] = dataset[i][j+1][k]
+                                else:
+                                    dataset[i][j][k] = dataset[i][j-1][k]
+                            else:
+                                dataset[i][j][k] = dataset[i][j][k-1]
+
+    dataset = (dataset - np.mean(dataset)) / np.std(dataset)
+
+    ans1 = np.isinf(dataset)
+    ans2 = np.isnan(dataset)
+
+    if not ((True in ans1) and (True in ans2)):
+        np.save('./data/sleepEDF-78/data_array/TF_data/TF_{}_mean_std.npy'.format(channel), dataset)
+
+
+if __name__ == '__main__':
+    path = Path()
+
+    fs = 100
+    overlap = 1
+    nfft = 256
+    win_size = 2
+
+    for channel in ['EEG_Fpz-Cz', 'EEG_Pz-Oz', 'EOG']:
+        print('-' * 15, 'Processing channel:{}'.format(channel), '-' * 15)
+        data_channel = data_array_concat(path_array=os.path.join(path.path_raw_data, channel))
+        X = np.zeros([data_channel.shape[0], 29, int(nfft / 2)])
+        print('Transform to TF images:')
+        for i in tqdm(range(data_channel.shape[0])):
+            Xi = spectrogram(data_channel[i, :], win_size * fs, overlap * fs, nfft)
+            Xi = 20 * np.log10(abs(Xi))
+            X[i, :, :] = Xi[:, 1:129]
+
+        print('Normalize:')
+        data_normalize(dataset=X, channel=channel)