ECG Neural Network Compression Practice: A Complete Workflow from Pruning & Quantization to Edge Deployment on ESP32

This article analyzes the compression scheme of an open-source ECG signal classification project, covering model training, pruning, INT8 quantization, and TensorFlow Lite conversion, and finally achieves efficient inference deployment on the ESP32 microcontroller. The project targets the MIT-BIH dataset, addresses resource constraints of edge devices, and provides an edge AI solution for arrhythmia detection.

Traditional ECG analysis relies on hospital equipment and manual interpretation, which cannot support continuous monitoring. The popularity of wearable devices has driven ECG analysis to the edge, but edge devices (such as MCUs) have constraints like small memory, limited performance, and power sensitivity. Core challenge: How to compress deep learning models while maintaining accuracy so they can run in real time on edge devices.

This project is open-sourced by alexToslev and provides a complete solution from training to deployment: train a classification model on the MIT-BIH dataset, convert it to TFLite format after pruning and quantization compression, and deploy it to ESP32. The MIT-BIH dataset contains over 100,000 ECG records labeled into 5 rhythm categories (N/V/A/F/Q), with clinical-level practical value.

The project adopts a three-layer compression strategy:

1. Lightweight CNN architecture: Use local receptive fields to capture ECG waveform features, and weight sharing to reduce the number of parameters;
2. Structured pruning: Remove redundant connections, reduce the number of convolution kernels, and lower computation and memory usage;
3. INT8 quantization: Through TensorFlow Lite post-training quantization, compress weights to 1/4 of the original size, and use SIMD to accelerate inference.

Training Phase: src/train_cnn.py processes MIT-BIH CSV data, splits it into 187-point heartbeat segments, and uses random translation/scaling for data augmentation; Compression Phase: src/compression/quantize_tflite.py loads the SavedModel, calibrates with a training subset, and converts it to an INT8 TFLite model; Deployment Phase: The TFLite model is deployed to ESP32 (dual-core 240MHz, 520KB SRAM), supporting TensorFlow Lite for Microcontrollers.

Performance of the compressed model:

• Size: Reduced from several MB to tens/hundreds of KB;
• Inference latency: Reaches millisecond level on ESP32;
• Accuracy: Classification accuracy drops by ≤2-3% compared to the floating-point baseline, meeting the needs of clinical real-time rhythm monitoring (focus on identifying dangerous arrhythmias).

Researchers: Provides a test benchmark for compression algorithms; Hardware developers: Can directly integrate the TFLite model to accelerate prototype development; the low cost of ESP32 (about $5) is conducive to large-scale deployment; Medical industry: Process sensitive ECG data locally to ensure privacy and security.

This project demonstrates the complete edge AI workflow: data preparation → model training → compression optimization → embedded deployment, proving the practical value of deep learning on resource-constrained MCUs. In the future, performance can be improved through more efficient architectures (such as temporal variants of MobileNet) and quantization-aware training, which is expected to realize professional-level rhythm monitoring AI built into smartwatches.