Building a Heartbeat Classifier in PyTorch Without nn.Linear or torch.optim

What you learn when you have to write every weight, loss and update yourself

Most PyTorch tutorials build a classifier in about ten lines. You stack a few nn.Linear layers, pick a loss from the library, hand everything to torch.optim and press go. It works, but it also hides almost everything that is actually happening.

For our Applied Machine Learning project, those shortcuts were banned. My teammate Moomal Eman and I had to build a network that classifies ECG heartbeats from the weight matrices up. This post walks through how we did it, the checks that told us it was working, and what the results revealed.

The task

An electrocardiogram (ECG) records the heart’s electrical activity, and each heartbeat produces a characteristic wave. We used the ECG Heartbeat Categorization Dataset from Kaggle, based on the MIT-BIH Arrhythmia Database.

Each heartbeat is a row of 187 numbers, measured 125 times per second and scaled between 0 and 1, plus a label for one of five classes: N (normal), S (supraventricular ectopic), V (ventricular ectopic), F (fusion) and Q (unknown or paced).

One example beat per class

In these examples, the normal beat stays low with one thin, tall spike. S and F look similar to it, with a sharp fall and a thin spike, while V and Q stay high through the middle of the beat. The flat line at the end of each one is zero padding, not heart activity.

We kept the provided test file of 21,892 beats aside, and split the training file into 78,799 training beats and 8,755 validation beats using a fixed random seed.

Step 1: Layers are just matrices

The architecture was fixed: 187 inputs, a hidden layer of 128 units, a hidden layer of 64 units and an output layer of 5 units. Each layer has a weight matrix and a bias vector, which we created ourselves as nn.Parameter tensors:

self.W1 = nn.Parameter(torch.randn(128, 187) * (2 / 187) ** 0.5)
self.b1 = nn.Parameter(torch.zeros(128))

The forward pass is three lines of matrix multiplication, with ReLU written as torch.clamp(z, min=0):

a1 = torch.clamp(x @ self.W1.T + self.b1, min=0)
a2 = torch.clamp(a1 @ self.W2.T + self.b2, min=0)
z3 = a2 @ self.W3.T + self.b3

Counting the parameters by hand is a good sanity check. Each layer has inputs × units weights plus one bias per unit:

Layer 1:  187 × 128 + 128 = 24,064
Layer 2:  128 × 64  + 64  =  8,256
Layer 3:   64 × 5   + 5   =    325
Total                     = 32,645

PyTorch counted the same 32,645.

Step 2: Why the starting weights matter

Notice the (2 / 187) ** 0.5 above. That is He initialisation, and it answers two questions.

Why not start at zero? If every weight is zero, every unit computes the same output, receives the same gradient and stays identical forever. We tested this on our network: every weight gradient came out exactly zero, so it could not learn from the input at all. Random weights break that symmetry.

Why √(2/n)? A unit adds up n weighted inputs, so the spread of its output grows with n. Dividing by n keeps the signal a steady size from layer to layer. The 2 makes up for ReLU, which zeroes out about half the values. Without this, the signal can fade away or blow up as it passes through the layers.

Step 3: A loss that does not overflow

For logits z and true class y, softmax cross-entropy can be written as −z_y + log Σ e^z. Computing e^z directly can overflow, so we used torch.logsumexp:

losses = -logits[torch.arange(len(y)), y] + torch.logsumexp(logits, dim=1)

A fresh network that knows nothing should give every class roughly equal probability, which means a loss near log 5 ≈ 1.61. Ours started at 1.7174, which is in that range for random starting weights.

Step 4: Trust, but verify the gradient

PyTorch computes gradients for us, but we checked one by hand. For the output layer, the gradient has a simple form: (1/B)(Ŷ − Y)ᵀA, where Ŷ holds the predicted probabilities, Y the one-hot labels and A the second hidden layer’s activations.

Get Tayyba’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

On the first 64 training rows, our hand calculation matched PyTorch’s stored gradients with a largest difference of 0.0, well within the required 10⁻⁵. As a test of the test, we also tried a deliberately wrong formula, and the difference jumped to about 0.9.

Step 5: Writing our own optimiser

Without torch.optim, the update rule is ours to write. For every parameter we kept a velocity, starting at zero:

g = p.grad + weight_decay * p   # weights only; biases skip the decay
v = beta * v + (1 - beta) * g
p -= alpha * v

Momentum (β = 0.9) averages recent gradients. When two consecutive mini-batch gradients point in opposite directions, they largely cancel. Starting from zero, a gradient of +1 gives a velocity of 0.1, and a following −1 brings it to just −0.01, so the weights don’t zig-zag.

Weight decay (λ = 0.0001) pulls each weight a little towards zero on every step. Even a weight that receives no gradient from the loss shrinks slowly, by roughly 3% per epoch with our settings.

We trained for 30 epochs with batches of 128 and tried three learning rates. Validation loss decided the winner:

Learning rate   Lowest validation loss   Validation accuracy
0.01            0.1618                   95.79%
0.1             0.0862                   97.64%
0.5             0.0722                   97.93%

Validation loss for each learning rate. The largest step size learned fastest and ended lowest

We chose 0.5 and saved the parameters from epoch 23, its lowest validation loss. Each run took under a minute on a laptop CPU.

Did it overfit?

Training and validation loss for learning rate 0.5

Only mildly. The two curves start together. Later, validation loss settles a little above training loss and stops improving after about epoch 20, while training loss keeps falling. At the best epoch, training and validation loss were 0.0560 and 0.0722, with accuracies of 98.26% and 97.93%. Saving the best epoch kept the final model clear of that late drift.

The results

Test accuracy          97.82%
Balanced accuracy      83.65%
Always-"N" accuracy    82.76%

Class                 N        S        V        F        Q
Per-class accuracy    99.60%   66.37%   91.37%   62.96%   97.95%

Test confusion matrix

97.82% sounds excellent, but a model that always says “normal” would already score 82.76%, because 18,118 of the 21,892 test beats are normal. Balanced accuracy, which counts each class equally, is 83.65%. Accuracy alone hides the weak spots.

The weakest class is fusion (F), at 62.96%. Of 162 fusion beats, 39 were predicted as normal and 21 as ventricular. A fusion beat is a mix of a normal and a ventricular beat, so it can resemble both. It is also very rare: 641 training rows against 72,471 normal ones.

If we continued, we would try a class-weighted loss, making mistakes on rare classes cost more. Weighting by rarity would give fusion a weight of about 27.3 and normal about 0.24. We did not run it, and it would probably cost some false alarms on normal beats.

A small front end

We wrapped the model in a Streamlit app. You can pick any test beat or a random beat of a chosen class, see it plotted, and see the predicted class, all five probabilities and whether the prediction was correct.

The front end in action

Project poster: the data, the hand-built network, our training choices and the results, all on one page

What building it by hand taught me

Writing every piece yourself makes the abstractions concrete. A layer is a matrix multiplication. An optimiser is three lines of arithmetic. A good starting point for the weights is a design decision with a reason behind it. And when the network reports 97.82%, you know enough about what is underneath to ask the next question: 97.82% of what?

Code, results and the app are on GitHub: [https://github.com/MoomalEmaan/Classifying-Heartbeats-from-ECG-Signals]