Imbalanced Heartbeat Classification Using CNN and Transformer

Yong Xia, Yueqi Xiong, Kuanquan Wang
Harbin Institute of Technology


Abstract

Aims: This study aimed to improve the inter-patient arrhythmia detection by using a novel deep model. Over the past few years, some so-called SOTA results didn't strictly follow the standard test protocol (test set and extra data can't be adopted in any way). In this paper, we will strictly obey the standard test protocol for a fair comparison. Methods: We proposed a novel framework by blending CNN and Transformer models. We firstly extracted the raw heartbeat signal around R peak and used a normalization strategy according to the locations of Waves R and T. Furthermore, we extract the RR features of heartbeat for cardiac rhythm representation. Then base on the extracted raw heartbeat signals and RR features, we used CNN to extract the local feature representation of a heartbeat signal. RR feature also reflects some local attributes, so the CNN output plus RR features were concatenated as the input of Transformer Encoder. For a continuous ECG signal, we constructed two vector series, namely single heartbeat wave serials and RR serials. They were input into the before and after of CNN Block respectively. The CNN network can learn the local feature representation of the raw heartbeat signal based the network optimization. Finally, each vector in the encoder output, corresponding to each raw heartbeat in the ECG recording, is mapped to a one-hot class code with a fully connected feed-forward network (FFN). Results: We considered all five heartbeat types mentioned in AAMI and followed the standard test protocol under inter-patient paradigm. Our model attains the SOTA for overall average accuracy of 97.66% on the MIT-BIH test set. Conclusion: Our model can boost the classification performance of imbalanced data set. Specially, the minority classes can achieve more balanced performance.