Rinil-Parmar/ecg-arrhythmia-kan-comparison

language: en license: mit library_name: pytorch tags:

• —ecg
• —arrhythmia
• —classification
• —pytorch
• —cnn
• —kan ---

ECG Arrhythmia Classification — CNN+MLP vs CNN+KAN

This repository provides two trained PyTorch models for ECG beat classification on the MIT-BIH Arrhythmia Database:

• —CNN+MLP (baseline)
• —CNN+KAN (proposed Kolmogorov–Arnold Network head)

Both models classify ECG beats into 5 AAMI-style classes and are evaluated on a held-out test set.

Models and Files

• —CNN+MLP weights: checkpoints/cnn_mlp.pth
• —CNN+KAN weights: checkpoints/cnn_kan.pth

Task

Time-series classification of single ECG beats.

Class Labels

Data

• —Dataset: MIT-BIH Arrhythmia Database (PhysioNet)
• —Lead: MLII (lead 0)
• —Sampling rate: 360 Hz
• —Beat window: 256 samples (128 before + 128 after annotation)
• —Normalization: per-beat z-score
• —Split: stratified train/val/test = 70% / 15% / 15%

Training

• —Optimizer: Adam
• —Learning rate: 1e-3
• —Weight decay: 1e-4
• —Batch size: 128
• —Epochs: up to 50, early stopping (patience 10)
• —Loss: class-weighted cross-entropy
• —Gradient clipping: max norm 1.0 (for KAN stability)

Results (Test Set)

Overall Metrics

Per-Class F1

How to Use (PyTorch)

import torch
from src.models.cnn import ECGCNN
from src.models.cnn_kan import ECGCNNWithKAN

# Choose model
model = ECGCNN(num_classes=5)      # or ECGCNNWithKAN(num_classes=5)

# Load weights
ckpt = torch.load(CHECKPOINT_PATH, map_location="cpu")
model.load_state_dict(ckpt["model_state"])
model.eval()

# Example input: [batch, 1, 256]
x = torch.randn(1, 1, 256)
proba = torch.softmax(model(x), dim=1)
pred = proba.argmax(dim=1).item()