README.md

---
library_name: pytorch
tags:
- ecg
- arrhythmia
- rhythm-classification
- ltaf
- physionet
- 1d-resnet
license: mit
datasets:
- physionet/ltafdb
---

# LTAF ECG Rhythm Classifier — RhythmResNet1D + TTA

A from-scratch 1D-ResNet trained on PhysioNet's
[Long-Term Atrial Fibrillation (LTAF)](https://physionet.org/content/ltafdb/)
database for **6-class rhythm classification** on two-lead 128 Hz ECG.

| Metric | Single-window | **+ 7-view TTA (recommended)** |
|---|---:|---:|
| Test accuracy | 0.636 | **0.684** |
| Test balanced accuracy | 0.740 | **0.778** |
| **Test macro F1** | **0.614** | **0.656** |

vs. frozen Chronos-2 + MLP baseline on the same 6-class subset:
test macro F1 = 0.299 — i.e. **+36 pp / 2.2× the F1**.

Per-class F1 (TTA-7): NSR 0.76, AFIB 0.62, SBR 0.82, AB 0.77, SVTA 0.15, B 0.82.

## Classes

| Code | Expansion |
|------|-----------|
| NSR  | Normal sinus rhythm |
| AFIB | Atrial fibrillation |
| SBR  | Sinus bradycardia (<60 bpm, sinus origin) |
| AB   | Atrial bigeminy (every other beat is an APC) |
| SVTA | Supraventricular tachyarrhythmia (≥3 consec SV ectopics @ >100 bpm) |
| B    | Ventricular bigeminy (every other beat is a PVC) |

`VT`, `T`, and `IVR` are excluded — their LTAF test supports (31, 26, 1) are too small for stable F1 estimation.

## Quickstart

```bash
pip install torch huggingface_hub numpy
```

```python
import numpy as np
import torch
from huggingface_hub import hf_hub_download
from model import RhythmResNet1D, RHYTHM_CLASS_NAMES

# Download checkpoint + model code from HF
ckpt = hf_hub_download("rmxjck/ltaf-ecg-rhythm-classifier", "best_classifier.pt")
model = RhythmResNet1D.load(ckpt, device="cuda")
model.eval()

# Input: (B, 2, 1280) — 10 s @ 128 Hz, 2 leads, per-channel z-scored.
x = torch.randn(1, 2, 1280).cuda()  # replace with real ECG
with torch.no_grad():
    logits = model(x)
    pred_idx = logits.argmax(-1).item()
print(model.class_names[pred_idx])
```

For best results, use the **7-view TTA** wrapper in `inference.py`
(averages softmax across 7 random window-start offsets — adds ~4 pp F1
at the cost of 7× inference compute).

```bash
python inference.py
```

## Architecture

`RhythmResNet1D(num_classes=6, n_channels=2, base_channels=64,
blocks_per_stage=2)`:

- **Stem:** Conv1d(2, 64, k=15, stride=2) → BN → ReLU → MaxPool(2).
- **4 ResNet stages × 2 basic blocks** (Conv1d k=7, BN, ReLU, Dropout, +skip).
  Channels: 64 → 128 → 256 → 512. Time downsamples 2× at the start of each
  stage past the first.
- **Head:** AdaptiveAvgPool1d → Linear(512 → 128) → ReLU → Dropout(0.2)
  → Linear(128 → 6).
- **Total parameters:** 8,794,246.

## Input format

- `(B, 2, 1280)` float32
- 2-lead ECG at **128 Hz** (LTAF leads `ECG1`, `ECG2`)
- 10 s window
- Per-channel z-scored: `(x - x.mean(axis=-1)) / x.std(axis=-1)`

## Test-time augmentation (TTA)

Pass a longer signal slice (≥1280 samples) to `predict_tta()` and it
samples 7 random 10 s windows, averages the softmax outputs, then
argmaxes. Why it helps: training uses random window-start sampling
within each rhythm bout, so the model learns to be invariant to that
shift. At eval time, taking multiple shifts and averaging cancels the
position-specific noise. **+4.2 pp test macro F1, no retraining.**

```python
# (2, 30*128) signal, 30 s long
cls, prob, full_probs = predict_tta(model, long_signal, n_views=7, device="cuda")
```

## Training recipe

```bash
.venv/bin/python scripts/train_ecg_rhythm_scratch.py \
    --arch resnet1d --window-sizes 10 \
    --epochs 30 --batch-size 64 --lr 5e-4 \
    --base-channels 64 \
    --use-val-as-train \
    --classes NSR AFIB SBR AB SVTA B \
    --output-dir results/ecg_classifier/sweep/c6_resnet1d_w10_e30_wide
```

- Dataset: LTAF train+val combined (75 records). 8 records held out for
  early stopping. Test (9 records, 3,716 windows) untouched.
- Loss: weighted cross-entropy with sqrt-dampened inverse-frequency
  class weights (cap 10), label smoothing 0.1.
- Cosine LR schedule from 5e-4 → 0 over 30 epochs. AdamW (wd 1e-4).
- Best checkpoint by held-out macro F1.
- Training time on a single H100 80GB: **~6 minutes**.

Source repo: `scripts/train_ecg_rhythm_scratch.py` and
`src/models/ts_llm/ecg_rhythm_scratch.py` in
[rmxjck/TSLM-Arena](https://github.com/rmxjck/TSLM-Arena).

## Test set details

LTAF held-out split (deterministic seed 42, record-level): 9 records
(`100, 104, 105, 11, 200, 32, 48, 49, 68`), 3,716 windows.

Confusion matrix (rows = true, cols = pred), with TTA:

|       | NSR | AFIB | SBR | AB  | SVTA | B   |
|-------|----:|-----:|----:|----:|-----:|----:|
| NSR   | 1109 | 286 | 95  | 114 | 185  | 35  |
| AFIB  | 189  | 628 | 25  | 29  | 294  | 14  |
| SBR   | 26   | 0   | 279 | 0   | 0    | 0   |
| AB    | 9    | 13  | 0   | 225 | 3    | 0   |
| SVTA  | 9    | 14  | 0   | 3   | 34   | 0   |
| B     | 4    | 0   | 0   | 1   | 3    | 90  |

Per-class supports: NSR 1824, AFIB 1179, SBR 305, AB 250, SVTA 60, B 98.

## What was tried and didn't help

This model was the best of 30+ experiments. What did *not* improve over
this baseline:

- HRV side-channel input (8-dim RR-derived features fused with CNN trunk):
  hurts F1 by 3-8 pp because the CNN already extracts equivalent
  information from raw QRS timing.
- Cross-corpus augmentation (MIT-BIH AFDB added to training): hurts
  AFIB F1 by 14 pp because AFDB's clean AFIB blocks bias the model
  toward over-calling AFIB on LTAF's paroxysmal transitions.
- Wider models (96-channel, 12 M params): overfits.
- Longer training (50 epochs): overfits.
- Multi-model soft-voting ensembles: members make correlated errors.
- Focal loss: matches CE within noise.
- Multi-scale training (5 / 10 / 30 s windows): underperforms 10 s alone.
- Bigger external models (torchecg ResNet-50 51.9 M, Stanford 27 M):
  underperform a 2.2 M home-rolled ResNet1D at 12 epochs.

## Not for clinical use

Research artifact only. **Not FDA-cleared.** Not suitable for triage,
diagnosis, or any patient-facing application. Uses the LTAF benchmark
which has known label noise from its original PhysioNet curation.

## Citation

```bibtex
@misc{petrutiu2008ltafdb,
  title         = {Abrupt Changes in Fibrillatory Wave Characteristics at the Termination of Paroxysmal Atrial Fibrillation in Humans},
  author        = {Petrutiu, Simona and Sahakian, Alan V. and Swiryn, Steven},
  year          = {2008},
  howpublished  = {PhysioNet},
  url           = {https://physionet.org/content/ltafdb/}
}
```