LEAST🫀

This repo covers the implementation of the following paper:

"Masked Autoencoders that Feel the Heart: Unveiling Simplicity Bias for ECG Analyses"

Requeirements

Python 3.12.4

Pytorch 2.6.0

pip install -r requirements.txt

Usage

For Pretraining

python main_pretrain.py --data_path "path/to/your/data" --batch_size "1024" --output_dir "path/to/your/outputdir" --log_dir "path/to/your/logdir" --model "LEAST" --mask_ratio "0.75"

For Full Fine-tuning

python main_finetune.py --model "LEAST" --finetune "path/to/your/weights" --data "path/to/your/data" --batch_size "512" --task cls  --classf_type multi_label 

The task variable can be set to one of three downstream tasks: 'cls', 'seg', or 'suv'. For example, when task='cls', available options for the data variable are: ptb-xl, cpsc_2018, chapman, and shaoxing.

For Linear Probing

python main_finetune.py --model "LEAST" --finetune "path/to/your/weights" --mode "linprobe" --data "path/to/your/data" --batch_size "512" --task cls  --classf_type multi_label 

Data preparation

For Pretraining

We constructed a large-scale multi-site ECG dataset consisting of 380,000 Chinese ECG records collected from over 300 clinical centers. In addition, we integrated the MIMIC-IV-ECG and CODE-15% datasets, resulting in a combined dataset containing approximately 1.53 million ECG records.

You can access MIMIC-IV-ECG and CODE-15% publicly. As for our own Chinese ECG dataset, it will be available upon request after the anonymization process is completed.

For Full Fine-tuning and Linear Probing

All datasets used for downstream tasks are publicly accessible.