gist-ailab / iitnet-official Goto Github PK

This is an official implementation for "Intra- and inter-epoch temporal context network (IITNet) using sub-epoch features for automatic sleep scoring on raw single-channel EEG".

License: MIT License

Python 78.96% Shell 21.04%

deep-learning deep-neural-networks biosignals eeg-signals sleep-staging eeg-classification pytorch

iitnet-official's Introduction

IITNet

By *Hogeon Seo, *Seunghyeok Back, *Seongju Lee, Deokhwan Park, Tae Kim, and Kyoobin Lee (*: Equal Contribution)

This repo is the official implementation of "Intra- and Inter-epoch Temporal Context Network (IITNet) using Sub-epoch Features for Automatic Sleep Scoring on Raw Single-channel EEG", Biomedical Signal Processing and Control 61 (2020): 102037. [paper]

The architecture of IITNet:

Updates & TODO Lists

(2022.03.22) IITNet official repository is released.
(2022.0X.XX) DeepSleepNet [1] baseline is added.

Getting Started

Environment

python >=3.7.0
pytorch >= 1.7.0 (or compatible version to your develop env)
numpy
scikit-learn
pandas
mne
terminaltables
termcolor

Install PyTorch with compatible version to your develop env from PyTorch official website.
Install remaining libraries using the following command.
```
pip install -r requirements.txt
```

Data Preparation

We evaluated our IITNet with MASS, SHHS, and Sleep-EDF dataset. You have to convert dataset from .edf into .npz format.

First, download .edf files and annotations MASS, SHHS, Sleep-EDF. For the MASS and SHHS dataset, you have to request for a permission to access their dataset. You can download SC subjects of Sleep-EDF using the following commands.

cd data
chmod +x download_physionet.sh
./download_physionet.sh

Second, preprocess .edf into .npz format using the following command. We referred prepare_physionet.py in DeepSleepNet repository.

python prepare_physionet.py

After preprocessing, the hierarchy of ./datasets/ directory will be the following:

./datasets/
├── MASS/
│   └── F4-EOG/
│       ├── 01-03-0001-PSG.npz
│       ├── ...
│       └── 01-03-0064-PSG.npz
├── SHHS/
│   └── C4-A1/
│       ├── shhs1-200001.npz
│       ├── ...
│       └── shhs1-205804.npz
└── Sleep-EDF/
    └── Fpz-Cz/
        ├── SC4001E0.npz
        ├── ...
        └── SC4192E0.npz

Each .npz file contains input eeg epochs with the shape of (total_num_epochs, 30 * sampling_rate) and target labels with the shape of (total_num_epochs) with the key 'x' and 'y', respectively.

Train & Evaluation

You can simply train and evaluate IITNet using just main.py.

$ python main.py --config $CONFIG_PATH --gpu $GPU_IDs

For evaluation only, add --test-only argument when you run the script

Example Commands

Train and Evaluation MASS (L=1) using single GPU (gpu_id=0)

$ python main.py --config ./configs/IITNet_MASS_SL-01.json --gpu 0

Train and Evaluation Sleep-EDF (L=10) using multiple GPUs (gpu_id=1,2)

$ python main.py --config ./configs/IITNet_Sleep-EDF_SL-10.json --gpu 1,2

Evalution trained SHHS (L=10) using single GPU (gpu_id=3)

$ python main.py --config ./configs/IITNet_SHHS_SL-10.json --gpu 3 --test-only

Outputs

For each fold, checkpoints that have the best validation loss are saved at ./checkpoints/CONFIG_NAME/.
Overall result is written in ./results/CONFIG_NAME.txt with the order of FOLD ACC MF1 KAPPA W-F1 N1-F1 N2-F1 N3-F1 REM-F1.
Overall and per-class results are printed like the following figures:

Citation

If you find this project useful, we would be grateful if you cite our work as follows:

@article{seo2020intra,
title={Intra-and inter-epoch temporal context network (IITNet) using sub-epoch features for automatic sleep scoring on raw single-channel EEG},
author={Seo, Hogeon and Back, Seunghyeok and Lee, Seongju and Park, Deokhwan and Kim, Tae and Lee, Kyoobin},
journal={Biomedical Signal Processing and Control},
volume={61},
pages={102037},
year={2020},
publisher={Elsevier}
}

Acknowledgement

This work was supported by the Institute of Integrated Technology (IIT) Research Project through a grant provided by Gwangju Institute of Science and Technology (GIST) in 2019 (Project Code: GK11470).

Reference

[1] A. Supratak, H. Dong, C. Wu, and Y. Guo, “DeepSleepNet: A model for automatic sleep stage scoring based on raw single-channel EEG,” IEEE Trans. Neural Syst. Rehabil. Eng., vol. 25, no. 11, pp. 1998–2008, 2017.

Licence

MIT licence

iitnet-official's People

Contributors

Stargazers

Watchers

Forkers

python-repository-hub bunkerdives wangliwen1994 resmedfeargalo

iitnet-official's Issues

Preprocessing of the MASS dataset

Hello,

I would like to use your program with the MASS dataset, to which I have access.
Unfortunately, I do not see how the file prepare_physionet.py could be called to parse said dataset.

The files in the dataset as I have it follow this format:
01-03-0001 Annotations.edf
01-03-0001 Base.edf
01-03-0001 PSG.edf
01-03-0002 Annotations.edf
... etc.

They are not properly parsed by prepare_physionet.py. For instance, they do not include any field for an electrode F4-EOG.

I could probably retro-engineer your process to an extent in order to make the preprocessing work, but as I wish to reproduce your results, I am weary of potentially making a mistake. (For instance, I could reproduce the F4-EOG signal from the available F4 and EOG signals, but as there are two EOG signals in MASS, I'd have to make an arbitrary choice...)

Could you please share the script you used to generate the preprocessed MASS data?

Thank you in advance.

refactored the code & detail

I have refactored the code, but it converges much more slowly than the original code. I don't know what details I might have missed.
`
def setProjectPath(projectPath):
import os
import sys
root=os.path.abspath(projectPath)
print(root)
sys.path.append(root)

os.environ["TORCH_HOME"]=r"E:\Data\torch-model"
os.environ["KMP_DUPLICATE_LIB_OK"]='TRUE'
import warnings  
warnings.filterwarnings("ignore", category=RuntimeWarning, module="mne")

projectPath=r'/mount/mount_project/test'
setProjectPath(projectPath)

iitnet_config_10={
"max_epochs": 500,
"dataset": "Sleep-EDF",
"signal_type": "Fpz-Cz",
"sampling_rate": 100,
"seq_len": 10,
"target_idx": -1,
"n_splits": 20,

"hidden_dim": 128,
"batch_size": 256,
"patience": 10,
"num_layers": 50,
"dropout_rate": 0.5,
"num_classes": 5,
"early_stopping_mode": "min",
"bidirectional": True,
"learning_rate": 0.005,
"weight_decay": 0.000001

}
iitnet_config_01={
"max_epochs": 500,
"dataset": "Sleep-EDF",
"signal_type": "Fpz-Cz",
"sampling_rate": 100,
"seq_len": 1,

"target_idx": -1,
"n_splits": 20,

"hidden_dim": 128,
"batch_size": 256,
"patience": 10,
"num_layers": 50,#resnet_18,34,50,101,152
"dropout_rate": 0.5,
"num_classes": 5,
"early_stopping_mode": "min",
"bidirectional": True,
"learning_rate": 0.005,
"weight_decay": 0.000001

}

from tqdm import tqdm

from project.dataset import loader
from torch.utils.data import *

train_dataset=loader.EEGDataLoader(iitnet_config_01,1)
train_dataloader=DataLoader(dataset=train_dataset,batch_size=128)

val_dataset=loader.EEGDataLoader(iitnet_config_01,1,mode='val')
val_dataloader=DataLoader(dataset=val_dataset,batch_size=128)
import project.model.iitnet.models.main_models as iitnet
import torch
model=iitnet.MainModel(config=iitnet_config_01)
model = torch.nn.DataParallel(model, device_ids=list(range(len('0'.split(",")))))
optimizer = torch.optim.Adam(model.parameters(),lr=iitnet_config_01['learning_rate'],weight_decay=iitnet_config_01['weight_decay'])
loss_fn=torch.nn.CrossEntropyLoss()
device=torch.device('cuda')
model.to(device)

for j in range(100):
model.train()
with tqdm(enumerate(train_dataloader),total=len(train_dataloader)) as t:
correct, total, train_loss = 0, 0, 0
for i,batch_data in t:

        x,y=batch_data
        x=x.to(device)
        y=y.to(device)
        total += y.size(0)
        y_hat=model(x)
        loss=loss_fn(y_hat,y)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_loss += loss.item()
        predicted=torch.argmax(y_hat,1)
        correct+=predicted.eq(y).sum().item()
        t.set_description_str(f'第{j}epoch')
        t.set_postfix_str('Loss: %.3f | TRAIN_Acc: %.3f%% (%d/%d)'
                        % (train_loss / (i + 1), 100. * correct / total, correct, total))
model.eval()
with tqdm(enumerate(val_dataloader),total=len(val_dataloader)) as t:
    correct, total, val_loss = 0, 0, 0
    for i,batch_data in t:
        x,y=batch_data
        x=x.to(device)
        y=y.to(device)
        total += y.size(0)
        y_hat=model(x)
        loss=loss_fn(y_hat,y)
        val_loss += loss.item()
        predicted=torch.argmax(y_hat,1)
        correct+=predicted.eq(y).sum().item()
        t.set_description_str(f'{j}epoch')
        t.set_postfix_str('Loss: %.3f | EVAL_Acc: %.3f%% (%d/%d)'
                        % (train_loss / (i + 1), 100. * correct / total, correct, total))