Classifying Heartbeats from ECG Signals with a Neural Network Built from Scratch

title: "Classifying Heartbeats from ECG Signals with a Neural Network Built from Scratch"
published: true
description: "A walkthrough of an ECG heartbeat classification project: data preparation, a hand-built PyTorch network, training, test results, and what the metrics reveal."

tags: machinelearning, pytorch, python, health

Classifying Heartbeats from ECG Signals with a Neural Network Built from Scratch

Note: This is an educational machine-learning project, not a medical device. Its predictions are not intended for diagnosis or treatment.

An electrocardiogram (ECG) records the electrical activity of the heart. In this project, I built a neural network to classify individual ECG heartbeat segments into five categories. The goal was not only to train a classifier, but also to understand the pieces behind it: preparing the data, implementing the network and loss, training with momentum, and evaluating the results beyond accuracy alone.

Project repository: Classifying Heartbeats from ECG Signals using Artificial Neural Networks

The classification task

Each example in the dataset contains 187 signal values and a class label. The network uses those 187 values to predict one of five heartbeat classes:

Label	Heartbeat category
N	Normal
S	Supraventricular ectopic
V	Ventricular ectopic
F	Fusion
Q	Unknown or paced

The project uses the MIT-BIH ECG Heartbeat Categorization dataset. It provides separate training and test CSV files, with 87,554 training records and 21,892 test records. The test data is held out for final evaluation. I split 10% of the training data into a validation set using a fixed random seed.

The classes are imbalanced: Normal beats are much more common than some of the other categories. That imbalance matters when interpreting model performance, as a high overall accuracy can hide weak results on less common classes.

Building the network

I implemented a fully connected network in PyTorch with this architecture:

187 inputs -> 128 hidden units -> 64 hidden units -> 5 output logits

The two hidden layers use ReLU activations. The weights use He initialization, which is designed to work well with ReLU, and the biases start at zero. The output layer produces logits for the five classes.

Rather than using PyTorch's ready-made linear layers for the network, the model defines its weight and bias parameters and matrix operations directly. The cross-entropy loss is also expressed using torch.logsumexp, which provides a numerically stable way to compute the softmax-based loss.

I also included a gradient check that compares a manually calculated output-layer gradient with the gradient produced by PyTorch autograd. This helps catch mistakes in the loss and gradient calculations before relying on the training loop.

Training and model selection

The model was trained for up to 30 epochs with mini-batches of 128 examples. The training loop uses momentum (0.9) and weight decay (0.0001, applied to weights). I compared three learning rates and selected the run with the lowest validation loss.

Learning rate	Lowest validation loss	Best epoch	Validation accuracy at best epoch
0.01	0.1618	30	95.77%
0.1	0.0857	28	97.59%
0.5	0.0729	23	97.78%

The 0.5 learning rate had the lowest validation loss, so its best checkpoint was saved and used for test evaluation.

Test results

On the held-out test set, the model correctly classified 21,399 of 21,892 examples:

Metric	Result
Test accuracy	97.75%
Balanced accuracy	83.76%

Balanced accuracy is the average of the per-class accuracies. It is useful here because each class contributes equally to that average, rather than allowing the largest class to dominate the metric.

Results by heartbeat class

Class	Accuracy
N — Normal	99.51%
S — Supraventricular	68.17%
V — Ventricular	90.95%
F — Fusion	62.35%
Q — Unknown/Paced	97.82%

The model performed very well on N, V, and Q, but its results were weaker on S and F. This gap explains why balanced accuracy is lower than the overall accuracy: the model’s performance is not equally strong across the five categories.

Confusion matrix

Rows are actual classes; columns are predicted classes.

Actual / predicted	N	S	V	F	Q
N	18,029	27	16	10	36
S	168	379	6	0	3
V	104	4	1,317	14	9
F	39	0	21	101	1
Q	29	2	4	0	1,573

Fusion (F) was the least accurately classified class. Among its errors, the most frequent prediction was N: 39 actual F examples were classified as normal. The project analysis also notes the small number of F examples in the training data, which gives the model fewer examples of that class's variation.

Why accuracy is not the whole story

An Always-N classifier would achieve about 82.76% accuracy on this test set because N is the majority class. The neural network's test accuracy is roughly 15 percentage points higher, showing that it learned to classify beyond simply predicting the majority category.

That comparison is helpful, but it does not mean every class is handled well. The S and F results show why class-level metrics and the confusion matrix should be reported alongside accuracy, particularly for imbalanced datasets.

What the results suggest

The experiments show that a small, manually implemented network can achieve strong overall performance on this dataset. At the same time, the per-class results point to a clear next step: improve recognition of the less accurately classified classes without sacrificing performance elsewhere.

One possible experiment is class-weighted cross-entropy, which would give errors on underrepresented classes more influence during training. This is a proposed direction, not a change evaluated in the current results. Any such change should be compared using the same held-out test set and assessed with balanced accuracy, per-class results, and the confusion matrix.

The training and validation curves are also useful for checking generalization. In this project, the recorded analysis describes a modest gap between training and validation loss, with validation loss fluctuating rather than rising continuously. That suggests some mild overfitting, but not a steadily worsening validation trend.

Exploring the project

The repository includes the training and evaluation code, saved model weights, result tables, plots, and a Streamlit application. The main implementation is organized as follows:

• src/load_data.py loads the dataset CSV files.
• src/prepare.py creates the train/validation split and tensors.
• src/model.py defines the network and cross-entropy loss.
• src/grad_check.py checks the output-layer gradient calculation.
• src/train.py trains models and compares learning rates.
• src/test_model.py evaluates the selected model and exports metrics.
• src/metrics.py computes accuracy, balanced accuracy, and the confusion matrix.
• src/plot_beats.py plots example beats for the five classes.
• app/app.py provides an interactive Streamlit prediction interface.

To reproduce the workflow, first download mitbih_train.csv and mitbih_test.csv from the MIT-BIH ECG Heartbeat Categorization dataset on Kaggle and place them in the repository's data/ directory. Install the project dependencies (PyTorch, NumPy, Pandas, Matplotlib, and Streamlit), then run these commands from the project root:

python -m src.grad_check
python -m src.train
python -m src.test_model
streamlit run app/app.py

The Streamlit app can classify a 187-value heartbeat uploaded as a CSV or, when the test CSV is present, a selected row from that file. It displays the predicted class, waveform, and model output probabilities.

Conclusion

This project was a practical exercise in implementing and evaluating a neural network, not just calling a training API. The final model reached 97.75% test accuracy, but its 83.76% balanced accuracy and lower S/F scores provide a more complete picture of what it learned.

The main lesson is to evaluate models in the context of their data. With imbalanced classes, one headline accuracy number is not enough: per-class performance and error patterns matter too. The repository contains the code and result artifacts for anyone who wants to inspect the implementation or try a follow-up experiment.

Sources

• Project repository
• MIT-BIH ECG Heartbeat Categorization dataset