Building an Image Classifier from Scratch

To get the most out of this blog post, I would recommend that you first read “Introduction to Neural Network”.

And here’s a link to the notebook with the all the code I used.

We have a few goals with this blog post:

1. Gain intuition on how a model predicts and trains.
2. A soft introduction to Softmax and Categorical Cross-Entropy functions.
3. Gain some intuition of the difference between a linear model vs a neural network.
4. Understand the math underneath the model.

At the end of the blog post, we would be able to create an image classifier from scratch that would identify a handwritten digit.

Intuition (The shouting crowd)

Imagine that you are inside a stadium. Inside the stadium, you see a big scoreboard that keep tracks of scores of ten teams: “Team 0”, “Team 1”, “Team 2”, up to “Team 9”.

Now imagine you have 784 different microphones: “Mic #1”, “Mic #2”, up to “Mic #784”. Each of the microphones has 10 knobs: “Knob 0”, “Knob 1”, “Knob 2”, “Knob 3”, up to “Knob 9”.

Every day, there will be a queue of 784 persons lined up outside of the stadium. Then one at a time, each person will go inside the stadium and you will give them a mic.

Now, imagine Person #1 goes inside the stadium. You will give the Mic #1 with knobs to Person #1. The knobs are fixed and Person #1 can’t adjust any of the 10 knobs in their microphone. Person #1 will then shout to Mic #1. After shouting, the microphone (after their voice passed through the knobs) sends a score to the scoreboard. The fixed knobs affects how much score each team gets.

For example, Mic #1 has the following knob settings:

# Imagine that hypothetical knobs has no numerical limits and can be set to any +/- numbers.
"Knob 0": 77,
"Knob 1": 5,
"Knob 2": 0,
"Knob 3": 10,
"Knob 4": 200,
"Knob 5": 70,
"Knob 6": 0,
"Knob 7": -100,
"Knob 8": 25,
"Knob 9": 130

When Person #1 shouts using Mic #1, it will send scores to the scoreboard amplified through the knobs. So for this particular instance it will send add a score to “Team 4” because “Knob 4” is set to be the highest with 200 value. But it will subtract to the existing score of “Team 7” because “Knob 7” is set to -100.

Now that Person #1 was done shouting and we already got the score. Person #2 will then enter the stadium and you will give them Mic #2 that has the knob settings:

"Knob 0": 1000,
"Knob 1": 0,
"Knob 2": 0,
"Knob 3": 0,
"Knob 4": 0,
"Knob 5": 0,
"Knob 6": 0,
"Knob 7": 0,
"Knob 8": 0,
"Knob 9": 0

What’s interesting is, each individual person also their own “loudness”-level to their voice. For example, Person #2’s voice is very quiet, with “loudness”-level = 0. In this case, even if Person #2’s Mic Knob 0 is set to very high, 1000, but because Person #2’s voice has no sound, it will not register any score to “Team 0”.

So you can imagine how a person’s voice loudness-level and their specific mic determines the score. Again, remember that each of their microphones has different fixed knob settings.

At the end of the day, all 784 persons shouted at 784 microphones and we’re left with a score board with various values. The “Team” with the highest score wins.

In this hypothetical scenario, we can think of these representations:

• 784 persons that enters the stadium and shouted are the 784 pixels (28×28) of the image.
• The loudness-level of each person represents the pixel intensity. A silent person represents a black pixel, a very loud person represents a white pixel.
• The team with the highest score in the scoreboard at the end of the day is our model’s prediction.

Supervised Learning

For this project, we will be using MNIST dataset. The MNIST dataset consists of 70,000 28×28 black-and-white images of handwritten digits. Here are a few samples of the images.

his dataset is labeled: we got pairs from the dataset, [image => label]. Which means what we are creating falls under “Supervised Learning”. To put it simple, this is how supervised learning works:

1. Labeled Data: You provide the algorithm with input-output pairs (e.g., images of fruits and their names).
2. Training: The model learns the mapping function from inputs to outputs by finding patterns in this data.
3. Prediction: Once trained, it can take new, unlabeled inputs (e.g., a new fruit picture) and predict the correct output (e.g., “apple”).

Preparing the data

Here’s how the images in the MNIST dataset looks

We will get these dataset from HuggingFace. Here’s the python code.

# Install deps
!pip install datasets

from datasets import load_dataset

# Load the full dataset (both train and test splits)
mnist_dataset = load_dataset("ylecun/mnist")

If you don’t understand the code above, it’s alright. Basically, we are just downloading the MNIST dataset and loading it so we can use it through the variable mnist_dataset.

Training and Test Split

If we look at mnist_dataset, we see this:

We see that the dataset is split into two parts: train and test. The train data set is as the name implies, this is the split we will use to train our Image Classifier model. Also, as the name implies, the test split is used to test and evaluate how well our model is when it sees a data that it never saw before (during training). It is important that we, who are creating the model, never look at the test set. I’ll dive deeper as to why this is, but the main reason is we don’t want to “fine tune” our model to the test set. Our model should do good on new data.

Now, let’s try to look at how a single training data point looks like.

mnist_dataset['train'][0]

As mentioned above, we will be using a labeled dataset. So each of our data points contains two pieces of information: image and label. The label is what the image is. For now, we are more interested as to what image looks like.

If we look at only the image, it is a PIL data.

It’s the image itself. But, we can’t do computation with images. We need numbers to do calculations. So we need to transform the image to an array of numbers.

import numpy as np

first_img_np = np.array(mnist_dataset['train'][0]['image'])
first_img_np

We can say that what you see above is how our computers see images. They are all numbers. As you see we have an array of 28 x 28, which is the size of the image. Each of the numbers in the array is pixel the image. The zeros (0) in the array are the black pixels, and 255 are the white pixels.

So as in our intuition analogy above, we want the crowd of 784 people (pixels) to queue in a line. Right now they are 28 x 28, we want them to be a single line.

first_img_vector = first_img_np.reshape(1, 784)

print(f"First image vector shape: {first_img_vector.shape}")

first_img_vector

Calculating the “shout” score

Now that our crowd is queued. We can let them in one-by-one and score their shout. The formula for the shout is this.

$z$: The score of the shout on each of the teams.
$w$: The fixed setting of the mic's knobs. This is our model's weights.
$x$: The intensity of the shout. This is our input. In this example, this is the pixel's whiteness intensity.
$b$: The bias.

We introduced $b$ above. We can think of the bias as helping our model learn the bias of the dataset. It is typical that a dataset has a “bias” of where most data points are. For example, if we look at test scores in a classroom, it is very likely that the score of the students are around the passing level. If the exam is very hard, then the scores would be around the failing mark. If the exam is very easy, then the scores would be around the perfect mark. Nonetheless, the scores are very likely to be around a “mean”.

Our image dataset (MNIST) is most likely to have a bias to a certain digit as well and we want to learn it so our model’s prediction will lean towards the bias.

Let’s build our model with what we know at this point.

import numpy as np

def loop_linear_model( single_image ):
	# For now we will process 1 image from the MNIST dataset.
	first_image = np.array( single_image )
	
	# Then line up the pixels.
	first_image_vector = first_image.reshape(1, 784)
	first_image_vector = first_image_vector[0]
	
	# Let's setup each of the mic's knobs.
	weights = np.random.rand(10, 784) # We have 10 knobs each in 784 mics.
	bias = np.zeros((10, 1)) # Since we are predicting a digit from 0 to 9. We use ten biases.
	
	# Setup our scoreboard. Team 0 - Team 9
	score_board = np.zeros((10, 1))
	
	# Let the pixel in one-by-one in the stadium
	for i in range(len(first_image_vector)):
		pixel = first_image_vector[i]
		
		# Loop through each of the knobs
		# Remember that each pixel adds a score to
		for j in range(10):
			# The specific knob in a specific mic.
			knob = weights[j, i]
			
			# The shout score per knob.
			shout_score = knob * pixel
			
			# Get the current score.
			current_score = score_board[j][0]
			
			# Let's add the scores in the scoreboard.
			score_board[j][0] = current_score + shout_score
			
	# Then let's loop through each of the scores in the scoreboard
	for j in range(10):
		# Get the current score.
		current_score = score_board[j][0]
		
		# Let's add the bias to each score.
		score_board[j][0] = current_score + bias[j][0]
	
	return score_board

score_board = loop_linear_model( mnist_dataset['train'][0]['image'] )

Running the code above will return something like this (remember that it will be different from yours if you run the code because we are initializing the weights with random values.)

Softmax function

What we got are raw scores. Remember that what we want as output for our model is a set of probabilities of our model’s prediction. For example, if we input an image of number “2” to our model, ideally the output would be something like this. [0, 0, 1, 0, 0, 0, 0, 0, 0]. This represents that our model is 100% confident to it’s prediction that the number is “2”. But, of course, models are not perfect. So we will most likely get something like this.

[
	0.01, // 1% confidence that it is "0"
	0.0, // 0% confidence that it is "1"
	0.7, // 70% confidence that it is "2"
	0.05, // 5% confidence that it is "3"
	0.05, // 5% confidence that it is "4"
	0.01, // 1% confidence that it is "5"
	0.0, // 0% confidence that it is "6"
	0.08, // 8% confidence that it is "7"
	0.04, // 4% confidence that it is "8"
	0.06 // 6% confidence that it is "9"
]

The sum of all the values should be 1 (100%).

Softmax is a mathematical function in machine learning that converts a vector of raw scores into a probability distribution, where each value is between 0 and 1, and all values sum up to 1, making them interpretable as probabilities for different classes in multi-class classification problems. It highlights the most likely class by giving it a higher probability while still assigning small probabilities to others, acting as a “soft” version of picking the maximum.

Mathematically, the Softmax function $\sigma(\mathbf{z})$ takes a vector $z$ of $K$ real numbers (logits) and maps them to a probability distribution of $K$ probabilities. In this case, we also call the Softmax function as $a$ (activation function) since it’s easier to write it as $$a$ compared to $\sigma(z_i)$.

For a vector $z=[z_1​,z_2​,…,z_K​]$, the Softmax of the $i$-th element is defined as:

Let’s see how our scores above will look like after passing it to Softmax function.

# For now, we will be using a loop so we can better visualize what's going on.
def loop_softmax(scores):
	softmax_probability_distribution = np.zeros(len(scores))
	# Get the max score.
	max_logit = np.max(scores)
	
	# Let's calculate the denominator first and only once.
	denominator = 0
	for j in range(len(scores)):
		z_j = scores[j][0] - max_logit
		denominator += np.exp(z_j)
		
	# Loop through each of the scores.
	for i in range(len(scores)):
		# Stability trick: subtract max value to prevent exponential overflow
		# This shifts values to be <= 0, preventing e^x from being huge.
		z_i = scores[i][0] - max_logit
		numerator = np.exp(z_i)
		
		softmax_score = numerator / denominator
		
	# Keep track of the softmax scores.
	softmax_probability_distribution[i] = softmax_score

	return softmax_probability_distribution

softmax_scores = loop_softmax(score_board)
for i in range(10):
	print(f"Score for team {i}: {softmax_scores[i]}")

Now we converted our raw scores to a probability distribution.

Categorical Cross-Entropy Loss Function

As always, we need to have a quantitative way to measure how “wrong” our model is. For categorical classification problems such as our image classifier, the commonly used loss function is the Categorical Cross-Entropy (CCE) function.

I’ve dived deeper on what Categorical Cross-Entropy (CCE) is here. I encouraged you to give it a read.

Formally, Categorical Cross-Entropy Cost Function measures the difference between two probability distributions: the true distribution $y$ and the predicted distribution $\hat{y}$.

For a single data sample $C$ possible classes, the loss $L$ is defined as:

Where:

• $C$ is the number of classes (e.g. 10 for MNIST digits).
• $y$ is the true label vector (One-Hot Encoded). If the true class is index $k$, then $y_k = 1$ and all other $y_j = 0$.
• $a_j$ is the predicted probability vector (output of our Softmax), where $a_j \in [0, 1]$ and $\sum a_j = 1$.

One-Hot Encoding

This is just a short review of what One-Hot Encoding is. As we all know, we need a mathematical representation of our categories. For example, our categories are.

[
	'Apple',
	'Orange',
	'Banana',
]

We represent (one-hot) them as:

[
	[1, 0, 0], # Apple
	[0, 1, 0], # Orange
	[0, 0, 1]  # Banana
]

We don’t represent them as 1 (Apple), 2 (Orange), 3 (Banana) because our model might think that Banana is somehow “better” or “greater” than an “Apple” because 3 > 1. But, as in most categorical problems, they are just the same, they are all fruits.

So in our MNIST. We represent the digits from 0 – 9 as:

[
	[1, 0, 0, 0, 0, 0, 0, 0, 0, 0], # 0
	[0, 1, 0, 0, 0, 0, 0, 0, 0, 0], # 1
	[0, 0, 1, 0, 0, 0, 0, 0, 0, 0], # 2
	[0, 0, 0, 1, 0, 0, 0, 0, 0, 0], # 3
	[0, 0, 0, 0, 1, 0, 0, 0, 0, 0], # 4
	[0, 0, 0, 0, 0, 1, 0, 0, 0, 0], # 5
	[0, 0, 0, 0, 0, 0, 1, 0, 0, 0], # 6
	[0, 0, 0, 0, 0, 0, 0, 1, 0, 0], # 7
	[0, 0, 0, 0, 0, 0, 0, 0, 9, 0], # 8
	[0, 0, 0, 0, 0, 0, 0, 0, 0, 1], # 9
]

Going back to the Categorical Cross-Entropy, because $y$ is One-Hot encoded (it has zeros everywhere except at the index of the true class, the target), the summation collapses. All terms where $y_j = 0$ becomes 0. Which leaves us to just:

This reveals the core mechanism: We only care about the log-probability the model assigned to the correct class.

• If $a_{target} \to 1$, then $\log(1) = 0$, so Loss $\to 0$.
• If $a_{target} \to 0$, then $\log(0) \to -\infty$, so Loss $\to \infty$.

Softmax Image Classifier Model

If you have worked on MNIST before and created an image classifier, you must know that it is not possible to create a good image classifier model just by using a linear model. I’ve written about this at length here.

But since the goal of this blog post is to help us have a good intuition and fundamentals of Machine Learning, I believe that showing the difference between just a linear model compared to a neural network would be extremely beneficial.

Backpropagation

As we know, training happens during backpropagation and we do this using Gradient Descent. And so we, need to compute for the gradients.

At the moment we have the following formulas.

To get the gradients, we compute using the following equations:

We know that:

Again, if you want to see how we got the formula above, you can read Categorical Cross-Entropy. With this we only need to get $\frac{\partial{z}}{\partial{w}}$ and $\frac{\partial{z}}{\partial{b}}$.

So putting them together, the formula to compute the gradients are:

Python code

class SoftmaxImageClassfier:
	
	def __init__(self):
		np.random.seed(42) # So we can reproduce.
		
		# Then we need to instantiate our weights (all the mics)
		# 10 knobs on each of the 784 mics.
		self.W = np.random.randn(10, 784) * 0.01 # Scaled down for stability
		self.b = np.zeros((10, 1)) # We have 10 biases.
		self.training_cost = []
	
	def forward(self, X):
		self.Z = np.dot(self.W, X.T) + self.b
		self.A = self.softmax(self.Z)
		return self.A
	
	def softmax(self, z):
		max_logit = np.max(z, axis=0, keepdims=True)
		capped_z = np.exp(z - max_logit)
		return capped_z / np.sum(capped_z, axis=0, keepdims=True)
	
	def train(self, X, y, learning_rate=0.01, epochs=50_000):
		m = X.shape[0] # number of samples.
		epsilon = 1e-8
		
		# We also need to one-hot encode our y.
		y_one_hot = np.zeros((10, m))
		
		for i in range(m):
			y_one_hot[y[i], i] = 1
		
		for epoch in range(epochs):
			self.forward(X)
			
			# Backprop
			error = self.A - y_one_hot
			d_w = (1 / m) * np.dot(error, X)
			d_b = (1 / m) * np.sum(error, axis=1, keepdims=True)
			
			# Gradient Descent
			self.W = self.W - (learning_rate * d_w)
			self.b = self.b - (learning_rate * d_b)
			
			if epoch % 100 == 0:
				cost = -np.sum(y_one_hot * np.log(self.A + epsilon)) / m # adding epsilon prevents log(0)
				self.training_cost.append((epoch, cost))
				print(f"Epoch: {epoch}, Cost: {cost}")

Matrix Math

Before we look how the SoftmaxImageClassfier() model trained and how well it does. If you have noticed instead of using loops to compute for the weights (shout score), we perform dot product. This is the same when we compute for the Softmax. Using matrix operations instead of looping makes our model more efficient and compute (train) drastically faster.

Softmax Model Training

A few notes regarding on how I trained the model.

1. We only trained using 20,000 training samples.
2. We only trained for 50,000 iterations.

As you can see at the end of training, the cost is 0.27391. Let’s set this aside for now and create and train a neural network with a hidden layer.

Neural Network with One Hidden Layer

Now let’s create a neural network model with one hidden layer.

class OneHiddenLayeredImageClassifier:

	def __init__(self):
		np.random.seed(42)
		
		# Layer 1: 784 Inputs -> 128 Hidden Neurons
		self.W1 = np.random.randn(128, 784) * 0.01
		self.b1 = np.zeros((128, 1))
		
		# Layer 2: 128 Hidden Neurons -> 10 Outputs
		self.W2 = np.random.randn(10, 128) * 0.01
		self.b2 = np.zeros((10, 1))
		
		self.training_cost = []
	
	def sigmoid(self, z):
		# np.clip prevents overflow/underflow errors in exp
		z_clipped = np.clip(z, -15, 15)
		return 1 / (1 + np.exp(-z_clipped))
	
	def sigmoid_derivative(self, a):
		return a * (1 - a)
	
	def softmax(self, z):
		max_logit = np.max(z, axis=0, keepdims=True)
		capped_z = np.exp(z - max_logit)
		return capped_z / np.sum(capped_z, axis=0, keepdims=True)
	
	def forward(self, X):
		# Layer 1 (Hidden Layer)
		self.Z1 = np.dot(self.W1, X.T) + self.b1
		self.A1 = self.sigmoid(self.Z1)
	
		# Layer 2 (Output Layer)
		self.Z2 = np.dot(self.W2, self.A1) + self.b2
		self.A2 = self.softmax(self.Z2)
		
		return self.A2
	
	def train(self, X, y, learning_rate=0.01, epochs=50_000):
		m = X.shape[0]
		epsilon = 1e-8
		
		# One-hot encoding
		y_one_hot = np.zeros((10, m))
		
		for i in range(m):
			y_one_hot[y[i], i] = 1
		
		for epoch in range(epochs):
			self.forward(X)
			
			# --- Backprop for Layer 2 (Output) ---
			dZ2 = self.A2 - y_one_hot
			dW2 = (1/m) * np.dot(dZ2, self.A1.T)
			db2 = (1/m) * np.sum(dZ2, axis=1, keepdims=True)
			
			# --- Backprop for Layer 1 (Hidden) ---
			# We multiply by W2 to see how much the hidden layer contributed to the error
			dA1 = np.dot(self.W2.T, dZ2)
			dZ1 = dA1 * self.sigmoid_derivative(self.A1)
			dW1 = (1/m) * np.dot(dZ1, X)
			db1 = (1/m) * np.sum(dZ1, axis=1, keepdims=True)
			
			# Gradient Descent Update
			self.W2 = self.W2 - (learning_rate * dW2)
			self.b2 = self.b2 - (learning_rate * db2)
			self.W1 = self.W1 - (learning_rate * dW1)
			self.b1 = self.b1 - (learning_rate * db1)
			
			if epoch % 1000 == 0:
				cost = -np.sum(y_one_hot * np.log(self.A2 + epsilon)) / m
				self.training_cost.append((epoch, cost))
				print(f"Epoch: {epoch}, Cost: {cost}")

Hidden Layer

Compared to the Softmax model above, we added a new (hidden) layer in this model between the input and the output. As to why we chose 128 neurons, there is no strict mathematical formula for choosing the size of a hidden layer, but here are a few guidelines that data scientists use for choosing:

1. The “In-Between” Rule: A general rule of thumb is to choose a number between the size of your input (784 pixels) and your output (10 digits).
2. Feature Compression: The goal of the hidden layer is to find patterns. By forcing 784 pixels to squeeze through 128 neurons, you force the network to drop the “noise” (like the blank black background) and only remember the important features (like loops and straight lines).
3. Computer Science Standards: We often choose powers of 2 (like 64, 128, 256, 512) because computers process these numbers slightly faster.

Sigmoid

You must have noticed that we used the Sigmoid function as the activation function in our first layer. The Sigmoid function squashes an input number ($z$) into the range (0,1).

Backpropagation

For our One-Hidden Layered Neural Network, we have the other formulas as the Softmax model plus the sigmoid function:

Getting the gradients would be a trickier because we now have a hidden layer. We will get the gradients on both layers.

Layer 2 (Output)

Layer 1 (Hidden)

Neural Network Model Training

At the end of the training our cost is 0.24965.

We can see better the difference between our two models during the training.

While our Softmax Model “learned” faster in the first part of the epoch, we can see that it plateau while our Neural Network with 1 Hidden Layer continues to learn.

Evaluating the models

And lastly, we want to evaluate our models. We want to see how well they do with data they have never seen before (the data from the test split).

# Load test data
X_test_raw = mnist_dataset['test']['image']
y_test = np.array(mnist_dataset['test']['label'])

# Preprocess test data (Flatten and scale)
# We follow the same preprocessing steps as the training data
X_test = np.array([np.array(img).reshape(-1) for img in X_test_raw]) / 255

def evaluate_model(model, X, y):
	# Forward pass to get probabilities
	# The model.forward() returns shape (10, m)
	probs = model.forward(X)
	
	# Get predictions: index of the max probability along axis 0
	preds = np.argmax(probs, axis=0)
	
	# Calculate accuracy
	accuracy = np.mean(preds == y)

return accuracy

# Calculate accuracy for both models
softmax_acc = evaluate_model(softmax_model, X_test, y_test)
nn_acc = evaluate_model(nn_model, X_test, y_test)

# Display results
print(f"Softmax Model Test Accuracy: {softmax_acc:.4f} ({softmax_acc*100:.2f}%)")
print(f"Neural Network Model Test Accuracy: {nn_acc:.4f} ({nn_acc*100:.2f}%)")

Softmax Model Test Accuracy: 0.9197 (91.97%)
Neural Network Model Test Accuracy: 0.9271 (92.71%)

It may look like that the difference is small but I want you to consider the following:

• It’s very likely that the Softmax Model hit it’s max accuracy.
• If we add more training samples and longer training, the Neural Network Model will only improve.
• Lastly, we are only dealing with simple dataset (grayscale images). If we deal with more complicated problem (animal classification) and dataset (HD picture of a dog), the gap between the two models will be much more apparent.

What’s next?

There are many things we can do to further improve our models and learn more concepts:

1. Use other activation function. In our Neural Model, we used Sigmoid. Try to use ReLu.
2. Our training is slow and we can’t add more training samples because we are training using the whole training dataset on each iteration. We can do mini-batch gradient descent instead.