Multi-Label Cardiac Arrhythmia Classification Using CNN with Self-Attention and LSTM

Zhaowei Zhu and Tingting Zhao
Ping An Technology


Abstract

Aims: In this study, we aimed to develop a multi-label classification model on 12-lead ECG signals, which can diagnose patients with zero, one or more types of cardiac arrhythmias out of atrial fibrillation, first-degree atrioventricular block, left bundle branch block, right bundle branch block, premature atrial complex, premature ventricular complex, ST-segment depression and ST-segment elevation. Methods: Data Preprocessing: The original ECG signal time series were divided into the frames by sliding windows with the length of 5 seconds and the stride of 1 second. The last frame of each lead was padded with 0 if the length is less than 5 seconds. Modelling: A convolutional neural network (CNN) module with self-attention mechanism was used to learn the feature representations of the frames, and these feature representations were fed into a long short-term memory (LSTM) module for the final prediction. We trained 1 model with all 12-lead ECG signals and 12 models with single-lead signal while filling the other leads with 0. All 13 models shared the same structure but were trained independently. The final probabilities of each class would be the weighted average predictions from the 13 models. The weights of each model were learned by a greedy algorithm. In addition, the traditional physiological signal processing method is applied to extract the clinical features. These features are concatenated with the hidden layer of DCNN, and enhance the classification performance. Results: On the released training data, our model achieved a score of 0.676 (geometric mean of F2: 0.783 and G2: 0.583) upon cross validation. On the hidden test set, the score was 0.669 (F2: 0.778, G2: 0.576). Our team name is HeartBeats. Conclusion: Our model can classify cardiac arrhythmias with good accuracy and generalizability.