Arrhythmia Detection Based on WGAN-GP and SE-ResNet1D

Abstract

A WGAN-GP-based ECG signal expansion and an SE-ResNet1D-based ECG classification method are proposed to address the problem of poor modeling results due to the imbalanced sample distribution of ECG data sets. The network architectures of WGAN-GP and SE-ResNet1D are designed according to the characteristics of ECG signals so that they can be better applied to the generation and classification of ECG signals. First, ECG data were generated using WGAN-GP on the MIT-BIH arrhythmia database to balance the dataset. Then, the experiments were performed using the AAMI category and inter-patient data partitioning principles, and classification experiments were performed using SE-ResNet1D on the imbalanced and balanced datasets, respectively, and compared with three networks, VGGNet, DenseNet and CNN+Bi-LSTM. The experimental results show that using WGAN-GP to balance the dataset can improve the accuracy and robustness of the model classification, and the proposed SE-ResNet1D outperforms the comparison model, with a precision of 95.80%, recall of 96.75% and an F1 measure of 96.27% on the balanced dataset. Our methods have the potential to be a useful diagnostic tool to assist cardiologists in the diagnosis of arrhythmias.