ai4healthuol/ecg-fm-benchmarking

This repository is a benchmarking and training framework for ECG-centric time-series foundation models, with some support for EEG and movement signals. The architecture is organized around a shared clinical time-series core, plus dataset-specific configuration and model wrappers for evaluating many pretrained checkpoints on many downstream tasks.

1. Top-level workflow

The repo has two main execution paths:

- Benchmark/evaluation via main_lite.py, main_lite_ecg.py, main_lite_eeg.py, main_lite_movement.py

- Pretraining via pretrain.py

These scripts sit on top of the clinical_ts package and are the main entry points for running experiments.

2. Data ingestion and preprocessing

Raw datasets are not stored in code; the repository expects preprocessed dataset artifacts produced externally, mainly through:

- preprocess_ecg_dataset.ipynb

- mimic_preprocessing.py

The data pipeline is driven by dataset-specific YAML configs under code/conf/data and code/clinical_ts/models/conf/data. The README also points to label-mapping spreadsheets and MIMIC-IV-ECG CSVs, which are part of the data preparation boundary.

3. Clinical time-series core package

code/clinical_ts is the central library. It contains the reusable building blocks shared across all tasks:

- data/ for dataset objects, transforms, and utilities

- models/ for supervised and foundation-model wrappers

- loss/ and metric/ for task objectives and evaluation

- task/ for task-specific orchestration

- utils/ for preprocessing, scheduling, evaluation, and helper functions

This is the main abstraction layer that separates dataset handling, model definition, and experiment execution.

4. Dataset layer

time_series_dataset.py, time_series_dataset_transforms.py, and time_series_dataset_utils.py define how signals and labels are loaded, normalized, sliced, and augmented. The dataset layer appears designed to support multiple modalities and multiple task types, with configuration rather than hard-coded dataset logic.

5. Task orchestration

code/clinical_ts/task/ecg.py and the main_lite_* scripts likely assemble the full experiment for a specific modality/task family. This layer is where a dataset, model, loss, and metrics are bound together into a runnable benchmark. It is the boundary between generic infrastructure and ECG-specific benchmark logic.

6. Model families

The repository supports both classic supervised architectures and foundation-model adapters:

- Lightweight supervised backbones such as basic_conv1d, inception1d, xresnet1d, s4_model, s42

- A common template_model.py / template_modules.py pattern for modular model construction

- Specialized ECG foundation model wrappers under models/ecg_foundation_models/

This separation suggests a plug-in style architecture where downstream tasks can swap encoders while keeping the rest of the pipeline stable.

7. Foundation model implementations

The ECG foundation-model area contains multiple model-specific subpackages:

- ecg_cpc

- ecg_jepa

- hubert_ecg

- st_mem

- merl

- ecgfm_ked

- ecg_founder

These packages each encapsulate their own backbone/config/utility code, indicating that pretrained checkpoints from different research groups are integrated through wrappers rather than a single unified backbone implementation.

8. Shared encoder/head abstractions

The clinical_ts/ts/ and clinical_ts/utils/heads.py modules indicate a decomposition into:

- an encoder/backbone that converts time series to latent representations

- a head that maps representations to classification/regression outputs

This is important for diagramming because the benchmark likely reuses the same task head interface across many model types and datasets.

9. Losses and metrics

loss/supervised.py and loss/selfsupervised.py split objective functions by training regime. metric/base.py and utility evaluators under utils/ provide task-specific evaluation for classification, regression, and clinical benchmark reporting. This supports both pretraining objectives and downstream fine-tuning/evaluation.

10. Configuration system

The repository uses many YAML configuration files rather than hard-coded experiment setup:

- global/run configs under code/conf

- model/data configs under code/clinical_ts/models/conf

- one data config per dataset/split/subject group for the benchmark

The architecture is configuration-driven, which is critical for reproducing the benchmark across 12 datasets and 26 tasks.

11. Experimental infrastructure

There are two native extensions under code/extensions/:

- cauchy/ with C++/CUDA code and benchmarking/tuning scripts

- kernels/

These are performance-oriented compute boundaries, likely used for accelerating specialized kernels in some models. They are optional infrastructure rather than core application logic, but they matter if you are drawing a diagram that includes training-time acceleration.

12. External assets and checkpoints

The benchmark depends on externally downloaded pretrained checkpoints for ECG-CPC, ECGFounder, ECG-JEPA, ST-MEM, MERL, ECGFM-KED, HuBERT-ECG, and ECG-FM. The repo is therefore not just training code; it is also a model-evaluation harness over third-party pretrained assets.

13. Output and evaluation boundary

The benchmark is focused on downstream performance reporting across many tasks, including regression and classification. The output boundary is therefore experiment results rather than an application UI or inference service. The presence of Confidence_Interval.xlsx suggests some postprocessing/statistical analysis outside the main training loop.

14. Runtime and tooling

The project is Python-based and appears designed for scripted batch execution, with sbatch run.sh in the README indicating HPC/Slurm usage. Environment files (env.yaml, ecg_fm_env.yaml) define the runtime dependencies, and the repo is organized for reproducible offline experimentation rather than interactive serving.

Overall, the architecture is: raw clinical time-series data -> preprocessing -> dataset/config loading -> shared time-series pipeline -> chosen model family/backbone -> task head -> loss/metrics -> benchmark results. The main architectural boundary is between dataset/task orchestration and the interchangeable model implementations, especially the multiple foundation-model adapters.