A Self-Contained STFT CNN for ECG Classification and Arrhythmia Detection at the Edge

Automated classification of Electrocardiogram (ECG) for arrhythmia monitoring is the core of cardiovascular disease diagnosis. Machine Learning (ML) is widely used for arrhythmia detection. The cloud-based inference is the prevailing deployment model of modern ML algorithms which does not always meet the availability and privacy requirements of ECG monitoring. Edge inference is an emerging alternative that addresses the concerns of latency, privacy, connectivity, and availability. However, edge deployment of ML models is challenging due to the demanding requirements of modern ML algorithms and the computation constraints of edge devices. In this work, we propose a lightweight self-contained short-time Fourier Transform (STFT) Convolutional Neural Network (CNN) model for ECG classification and arrhythmia detection in real-time at the edge. We provide a clear interpretation of the convolutional layer as a Finite Impulse Response (FIR) filter and exploit this interpretation to develop an STFT-based 1D convolutional (Conv1D) layer to extract the spectrogram of the input ECG signal. The Conv1D output feature maps are reshaped into a 2D heatmap image and fed to a 2D convolutional (Conv2D) neural network (CNN) for classification. The MIT-BIH arrhythmia database is used for model training and testing. Four model variants are trained and tested on a cloud machine and then optimized for edge computing on a raspberry-pi device. Weight quantization and pruning techniques are applied to optimize the developed models for edge inference. The proposed classifier can achieve up to 99.1% classification accuracy and 95% F1-score at the edge with a maximum model size of 90 KB, an average inference time of 9 ms, and a maximum memory usage of 12 MB. The achieved results of the proposed classifier enable its deployment on a wide range of edge devices for arrhythmia monitoring.