An ECG foundation model for generalizable cardiac function prediction across the lifespan

Training and fine-tuning framework for ECG foundation models.

Overview

This repository provides tools for:

• Pretraining ECG foundation models with Fyler-code as supervision
• Fine-tuning pretrained models on downstream tasks
• Evaluation on various ECG-based prediction tasks

Setup

# Create environment
conda env create -f environment.yml
conda activate ECG-FM

# Install dependencies
pip install -r requirements.txt

Quick Start

Pretraining

sbatch pretrain_fc.sh              # Pretraining with Fyler codes

Fine-tuning

sbatch finetune_fc.sh              # Fine-tune Fyler-code pretrained model

# Other baseline ECG foundation models: 
sbatch finetune_ecgfounder.sh      # Fine-tune ECGFounder model
sbatch finetune_vitmae.sh          # Fine-tune ViT-MAE model

Training from scratch

sbatch train_scratch.sh            # Train a classification model from scratch (ResNet)

Extract Embeddings

The Fyler-code pretrained checkpoint from scripts/pretrain_fc.sh can be loaded with fyler_embedding.py. This script reconstructs the pretrained ResNet encoder and returns the feature vector right before the final classification layer.

For the current pretrained model in ../ECG_Fyler, the expected ECG input shape is [batch, 12, 2048], and the output embedding dimension is 2560.

Python API:

import torch
from fyler_embedding import load_fyler_encoder

model = load_fyler_encoder("../ECG_Fyler/ECG_Fyler.ckpt", device="cpu")

# ECG shape: [batch, 12, 2048]
ecg = torch.randn(1, 12, 2048)
embedding = model.encode(ecg)
print(embedding.shape)  # [1, 2560]

Command line:

python fyler_embedding.py \
  --checkpoint ../ECG_Fyler/ECG_Fyler.ckpt \
  --input your_ecg.npy \
  --output your_embedding.npy

Supported .npy ECG input shapes:

• [12, 2048]
• [2048, 12]
• [batch, 12, 2048]
• [batch, 2048, 12]

Project Structure

├── train_model_fc.py              # Pretraining and finetuning script 
├── train_model_spec.py            # Finetuning ViTMAE (Spectrogram-input) 
├── data.py                        # Data utilities
├── preprocessing/                 # Data preprocessing tools
├── postprocessing/                # Results postprocessing
├── models/                        # Model architectures
├── fyler_embedding.py             # Load pretrained checkpoint and extract embeddings