Rethinking ECG Classification with Neural Networks as a Sequence-to-Sequence Task

Philipp Sodmann1, Lars Kaderali1, Marcus Vollmer2
1Universit├Ąt Greifswald, 2Institute of Bioinformatics, University Medicine Greifswald


Introduction: The ECG is an important non-invasive tool that allows for a rapid diagnosis of life-threatening diseases. Correct interpretation is a skill that requires a lot of practice to learn and master. To aid routine diagnosis, we have trained a network that can segment the ECG and uses the information to identify cardiac pathologies.

Methods: We trained a fully convolutional U-Net with multiple heads to solve problems that build on each other. The segmentation head is trained to segment the ECG into P waves, QRS complexes, and T waves. The head for beat classification is trained to classify a QRS complex as normal, ventricular or supraventricular extrasystole. The classification head was trained to detect pathologies in the ECG. For that we randomly sampled 1500 12-lead ECGs, stratified over all pathologies and manually segmented the ECG and classified each heartbeat. Then, we trained the neural network by optimizing for three distinct criteria and summing their individual losses in a weighted manner. We used Tversky loss for ECG segmentation and binary cross-entropy for beat and pathology classification. The network was trained on 10-second excerpts of ECGs. For ECGs without a manual segmentation, we set the loss for segmentation and beat classification to zero to enable semi-supervised training. Pathologies that were never annotated in a database, like sinus tachycardia in the CPSC dataset, were masked for records taken from that particular database. Generalizability was tested in participation of the PhysioNet/CinC Challenge 2021. We trained a single model and used zeroing of leads for augmentation to solve the classification tasks on 6-, 3- and 2-lead ECGs.

Results: Our team HeartlyAI achieved a challenge score of 0.52 on 12-lead ECGs, and 0.48, 0.52, 0.51 on six, three and two leads respectively. The challenge metrics on a 10% test split were 0.72, 0.58, 0.57 and 0.44.