Tumo505/SSL-ECG-Classification-model-card

ECG classification: augmentation controls and encoder adaptation

This release contains the fully fine-tuned domain-inspired SimCLR and BYOL classifiers from the completed patient-separated PTB-XL study: ten seeds per framework, twenty distinct classifiers. The repository root duplicates SimCLR seed 42 for a stable default. Seed 42 is the first fixed protocol seed, not the best test seed. These are separate classifiers, not a trained ensemble.

The scientific contribution is controlled evaluation of established SSL methods under generic/domain policies, gate/RMS matching and encoder adaptation, supported by a reproducible seven-transformation library. No new SSL objective is claimed.

Input and output

• Input: float ECG array (batch, 12, 1000), ten seconds at 100 Hz.
• Lead order: I, II, III, aVR, aVL, aVF, V1, V2, V3, V4, V5, V6.
• Standardize each record and lead: (x - mean) / max(population SD, 1e-6).
• Model forward expects already standardized tensors. inference.py performs this preparation for raw arrays. Do not silently pad a shorter recording or interpret 5,000 samples at 100 Hz as a ten-second recording. Resample upstream if needed.
• Output: five logits in order NORM, CD, HYP, MI, STTC. Apply sigmoid independently; labels can co-occur. The primary decision threshold is 0.5, not softmax/argmax.
• Sigmoid outputs are model scores; clinical calibration is not established.

Quick start

Install requirements.txt. After upload, AutoModel resolves the custom code:

import numpy as np
import torch
from transformers import AutoModel

repo = "Tumo505/SSL-ECG-Classification-model-card"
model = AutoModel.from_pretrained(repo, trust_remote_code=True).eval()
raw = np.load("ecg_100hz.npy", allow_pickle=False).astype(np.float32)
assert raw.shape == (12, 1000)  # standard lead order, ten seconds
x = (raw - raw.mean(-1, keepdims=True)) / np.maximum(raw.std(-1, keepdims=True), 1e-6)
with torch.inference_mode():
    probabilities = model(input_values=torch.from_numpy(x[None])).logits.sigmoid()[0]
classes = ["NORM", "CD", "HYP", "MI", "STTC"]
print(dict(zip(classes, probabilities.tolist())))
print([name for name, score in zip(classes, probabilities) if score >= 0.5])

Inspect the custom Python files before enabling trust_remote_code. Pin a Hub commit for reproducible use. For BYOL seed 42, pass subfolder="models/byol_domain_finetuned/seed_42" to from_pretrained. All ten seeds 42, 52, 62, 72, 82, 92, 102, 112, 122, 132 are available for each framework. For a local checkout, python inference.py ecg_100hz.npy --model . performs normalization and returns multilabel scores. No hosted inference service is claimed.

Architecture and training

Six Conv1D/BatchNorm/ReLU blocks: channels 64,64,128,128,256,256; kernels 11,7,7,5,5,3; strides 2,1,2,1,2,1. Global average pooling yields 256 features; a linear head produces five logits. SSL projector/predictor heads are absent from these downstream weights. Each model.safetensors contains the complete fine-tuned encoder and classification head, including batch-normalization statistics.

SSL: 20 epochs, batch 128, Adam 0.001, cosine annealing. SimCLR projector 256->2048->128, temperature 0.07; BYOL projector 256->2048->256 and EMA 0.999. The domain policy uses jitter/scaling plus independently gated warping, mixup, bandpass variation, CutMix, motion artefacts, channel noise and temporal dropout. Clinical realism of these transformations is not experimentally established. Downstream: full encoder/head training for 20 epochs, batch 64, Adam 0.001, focal loss gamma 2, positive/negative alpha 0.25/0.75, weighted record sampling. The saved epoch maximizes validation macro-F1 at 0.5; test results never select the checkpoint or default release seed. No retraining was performed for export.

Dataset and evaluation

PTB-XL v1.0.3: 21,799 records, 18,869 patients. Recommended patient-disjoint folds 1-8 train, 9 validation, 10 test. Pretraining uses 17,418 training records only. Eligible diagnostic counts: 17,084 train, 2,146 validation, 2,158 test. The 10% labelled training budget is 1,708; conditions share subsets within seeds.

Ten-seed means +/- sample SD (primary threshold 0.5):

Fully trained method	Macro-F1	Macro-AUROC
Supervised	0.5802 +/- 0.0095	0.8534 +/- 0.0064
Domain SimCLR	0.6061 +/- 0.0196	0.8572 +/- 0.0072
Domain BYOL	0.6067 +/- 0.0138	0.8594 +/- 0.0078
Gate/RMS matched generic SimCLR	0.5958 +/- 0.0189	0.8570 +/- 0.0049
Gate/RMS matched generic BYOL	0.5981 +/- 0.0074	0.8568 +/- 0.0063

These are study averages, NOT the scores of the repository-root checkpoint. Root SimCLR seed 42 alone: F1 0.6181, AUROC 0.8651, selected epoch 18. Each model's provenance.json contains its own exact metrics, validation-selected epoch, optional validation thresholds, fingerprints and parent checkpoint hash. inference.py uses 0.5; it does not apply the secondary validation thresholds.

Domain fine-tuning F1 gains over supervised survive the dedicated family (paired Holm p=0.0195), but not the broader label-efficiency family. Neither framework establishes an AUROC gain over supervised or a fine-tuned gain over matched generic controls. Frozen domain encoders underperform supervised. Gate/RMS matching weakens augmentation-specific attribution, and omission effects are mixed. In the broader 570-run study, no declared label-efficiency or MIT-BIH contrast survives its panel correction. MIT-BIH is a binary ten-second window task on one previously evaluated fixed subject partition across ten training seeds, not an AAMI beat benchmark or untouched external confirmation. This release exports PTB-XL models, not the separately trained MIT-BIH classifiers.

Intended use and limitations

Research and reproducibility only; not prospectively validated for diagnosis, screening or treatment decisions. The five multilabel superclasses are not exhaustive diagnoses. Full twelve-lead input and correct lead order are required. Performance does not establish robustness across populations, devices or sites. Gate/RMS matching does not establish morphological, spectral or clinical fidelity. Omission studies do not identify augmentation interactions. Equal epochs are not equal compute budgets. Follow-up contrasts were declared after original outcomes were known; post-hoc calibration/policy contrasts remain post hoc. Confidence intervals are pointwise, conditional on the observed training/validation cohorts. Failure to reject a difference is not equivalence. Ten-seed mean results do not guarantee the performance of any one model or an unevaluated ensemble.

Provenance and availability

Safetensors exports preserve the original model tensors; release validation checks them against source checkpoints and saved held-out predictions. No raw waveforms, patient IDs, record-level predictions, optimizer states or credentials are included. release.json indexes the twenty models; study_metrics.json contains aggregate scientific evidence; validation_report.json and release_manifest.json record verification and file checksums. The release deliberately changes the legacy class order, input contract and architecture; callers must use this config and model code together rather than mixing them with older cached files.

Code and model release: Apache-2.0; see LICENSE. PTB-XL and MIT-BIH retain their own dataset terms; no dataset is redistributed here.

Source: https://github.com/Tumo505/SSL-for-ECG-classification PTB-XL: https://physionet.org/content/ptb-xl/1.0.3/ (Wagner et al., 2020, https://doi.org/10.1038/s41597-020-0495-6). MIT-BIH: https://physionet.org/content/mitdb/1.0.0/ (Moody and Mark, 2001, https://doi.org/10.1109/51.932724). SimCLR: Chen et al., 2020, https://proceedings.mlr.press/v119/chen20j.html. BYOL: Grill et al., 2020, https://arxiv.org/abs/2006.07733. CLOCS: Kiyasseh et al., 2021, https://proceedings.mlr.press/v139/kiyasseh21a.html.