ShapeEFM (Shape-based ECG Foundation Model) is a two-stage self-supervised learning framework for ECG analysis. It learns rich ECG representations through shape tokenization and masked language modeling, enabling effective transfer learning for downstream ECG classification tasks.
- Shape-aware tokenization: A VQ-VAE tokenizer learns a discrete codebook of 8192 ECG shape primitives via Perceptually Important Points (PIP) extraction
- Masked language pretraining: A 12-layer Transformer predicts masked shape tokens, learning contextual ECG representations
- Data-efficient transfer: Linear probing on downstream tasks with as little as 1% labeled data
- Multi-dataset support: MIMIC-IV ECG, EchoNext, and CSN arrhythmia
cd ShapeEFM
pip install -r requirements.txtTrain the VQ-VAE tokenizer to learn discrete ECG shape primitives:
bash tokenizer.shTrain the masked language model on tokenized ECG
bash train.shLinear probing freezes all layers of the pre-trained model except the final classification head, and only trains this linear layer on the downstream task.
| Parameter | Value | Description |
|---|---|---|
seq_len |
248 | Number of ECG shape segments per sample |
embed_dim |
256 | Transformer embedding dimension |
codebook_size |
8192 | Number of discrete shape tokens |
codebook_dim |
128 | Codebook embedding dimension |
mask_ratio |
0.5 | Fraction of tokens masked during pretraining |
sampling_rate |
500 Hz | ECG sampling rate |
See requirements.txt for the full list. Key dependencies:
- PyTorch >= 2.4.0
- timm >= 1.0.24
- numpy, pandas, scikit-learn
- wfdb (for ECG waveform processing)
- pyhealth (for healthcare data utilities)