5-Stage Pipeline in <15 Lines

The same pattern works for any task — swap the dataset and task class to move between mortality, readmission, drug recommendation, or imaging.

from pyhealth.datasets import MIMIC3Dataset
from pyhealth.tasks import MortalityPredictionMIMIC3
from pyhealth.datasets import split_by_patient, get_dataloader
from pyhealth.models import Transformer
from pyhealth.trainer import Trainer

if __name__ == "__main__":
    # 1. Load data
    dataset = MIMIC3Dataset(root="data/", tables=["DIAGNOSES_ICD", "PROCEDURES_ICD"])
    samples = dataset.set_task(MortalityPredictionMIMIC3())

    # 2. Split & load
    train_ds, val_ds, test_ds = split_by_patient(samples, [0.8, 0.1, 0.1])
    train_loader = get_dataloader(train_ds, batch_size=32, shuffle=True)

    # 3. Train
    model = Transformer(dataset=samples)
    trainer = Trainer(model=model)
    trainer.train(train_loader, val_loader, epochs=50, monitor="pr_auc")
    trainer.evaluate(test_loader)

Built for Scale — Efficient by Design

Benchmarked on MIMIC-IV at 4 parallel workers. PyHealth 2.0 uses a memory-mapped architecture that dynamically adapts to your hardware — whether you have 2 cores or 64, it scales without manual memory management so you can focus on the ML, not the infrastructure. Pandas values shown at 1 worker only (†) as it could not scale beyond single-threaded execution. PyHealth 2.0 even runs the full in-hospital mortality task on a consumer laptop in roughly 5 GB of memory.

† Pandas measured at 1 worker; failed beyond single-threaded execution. All other methods at 4 workers on MIMIC-IV. Laptop rows are PyHealth 2.0 on a MacBook Pro M2 Pro and an ASUS Zephyrus G16 at 4 workers, benchmarked for in-hospital mortality. Full benchmark in the PyHealth 2.0 paper.

How PyHealth Standardizes Clinical AI Development

A unified API from raw clinical data to trustworthy, interpretable models — follow the pipeline step by step.

Resources

Everything you need to get started or get involved.