DeepECG-Bench

Train and test a model

from deepecgkit import ECGDataModule, ECGTrainer, build_model

ECGTrainer.seed_everything(42)

# PhysioNet 2017 (about 100 MB) as 10 s single-lead segments, resampled from
# 300 Hz to the 500 Hz the benchmark uses. The 70/20/10 split is by recording
# and stratified by class, so no recording has segments in both train and test.
data = ECGDataModule("af-classification", sampling_rate=500, download=True)

# Any name from `deepecg list-models` can replace "kanres".
model = build_model("kanres", data.input_channels, data.num_classes)

# Training ends by restoring the epoch with the lowest validation loss.
trainer = ECGTrainer(model)
trainer.fit(data, epochs=50, checkpoint_dir="runs/kanres-af")
print(trainer.test(data))  # {"test_loss": 0.806, "test_acc": 0.655}

The figures in this Quick Start come from the benchmark's KanRes run with this setup (seed 42, best epoch 50 of 50), which also scores macro F1 0.453, MCC 0.317 and AUROC 0.795 on the same test split. A new run differs slightly with hardware and library version.

The trainer reads the task from the data: multi-label datasets such as PTB-XL train with a sigmoid per class, with no setting to change. For a training loop of your own, data.train_dataloader() yields batches of signals shaped (32, 1, 5000) with their labels.

Predict on a raw recording

from deepecgkit import predict

# The checkpoint records its architecture, class names and preprocessing, so a
# recording at another rate and length is resampled, cut and normalised to match.
result = predict("ecg.npy", checkpoint="runs/kanres-af", sampling_rate=250)
print(result.label, result.confidence)
print(result.probabilities)  # every class, averaged over the recording's segments

A run directory stands for its best checkpoint wherever a checkpoint is expected. To predict many recordings, load the model once with ECGPredictor.from_checkpoint("runs/kanres-af") and call its predict method for each.

Explain the predictions

from deepecgkit import explain

# Saliency for every test segment, averaged over R-peak-aligned beats. The
# checkpoint does not record the rate the beat windows assume, so state it.
result = explain(data, checkpoint="runs/kanres-af", sampling_rate=500)

# Share of the average beat's attribution in each region
result.pqrst_fractions   # P 0.315, QRS 0.082, T 0.343, baseline 0.260
# Each share over the region's share of the beat, so 1 is chance
result.pqrst_enrichment  # P 1.26, QRS 0.66, T 0.92, baseline 1.04

On multi-lead data such as PTB-XL, result.lead_importance also gives each lead's share of the attribution.