This repository provides two trained PyTorch models for ECG beat classification on the MIT-BIH Arrhythmia Database:
Both models classify ECG beats into 5 AAMI-style classes and are evaluated on a held-out test set.
1import torch
2from src.models.cnn import ECGCNN
3from src.models.cnn_kan import ECGCNNWithKAN
4
5# Choose model
6model = ECGCNN(num_classes=5) # or ECGCNNWithKAN(num_classes=5)
7
8# Load weights
9ckpt = torch.load(CHECKPOINT_PATH, map_location="cpu")
10model.load_state_dict(ckpt["model_state"])
11model.eval()
12
13# Example input: [batch, 1, 256]
14x = torch.randn(1, 1, 256)
15proba = torch.softmax(model(x), dim=1)
16pred = proba.argmax(dim=1).item()