Building a Robust Handwritten Digit Classifier with TensorFlow and Keras

From preprocessing and architecture design to regularization and evaluation — a complete deep learning pipeline on MNIST

If there is one project that teaches the foundations of deep learning surprisingly well, it is handwritten digit classification.

The MNIST dataset may look simple at first: 70,000 small grayscale images containing handwritten digits from 0 to 9. But building a model that performs reliably on unseen handwriting introduces many of the same engineering questions that appear in much larger machine learning systems:

• How should the input data be prepared?
• How should a neural network architecture be designed?
• How do I prevent overfitting?
• How should I know whether the model is actually learning?
• What should I do when training stops improving?
• Is accuracy alone enough to evaluate a classifier?
• How would I move a notebook experiment toward production?

I explored these questions while building a handwritten digit classifier during my AI/ML Engineering Bootcamp at Ostad (Batch 6).

Rather than treating MNIST as a simple “train a model and get an accuracy score” exercise, I wanted to understand the reasoning behind each stage of the pipeline.

The result was a fully connected Deep Neural Network (DNN) using TensorFlow and Keras, with:

• Input normalization
• Flattening
• One-hot encoded labels
• Three Dense hidden layers
• Batch Normalization
• ReLU activation
• Dropout regularization
• Adam optimization
• EarlyStopping
• ReduceLROnPlateau
• Confusion matrix analysis
• Confidence-based predictions

The final model achieved approximately 98.5% test accuracy on the MNIST test set.

But the accuracy was only one part of the learning experience.

The more valuable outcome was understanding why the model worked, where it failed, and how I could improve it for a more realistic production environment.

1. Understanding the MNIST Dataset

MNIST stands for Modified National Institute of Standards and Technology.

It contains:

• 60,000 training images
• 10,000 test images
• Images of size 28 × 28 pixels
• Grayscale pixel values from 0 to 255
• Ten classes representing digits 0 through 9

Each image represents a handwritten digit.

For example:

0
1
2
3
...
9

The original images are small, but the underlying problem is meaningful: the model must learn to distinguish different handwriting styles.

The same digit can appear with different:

• Stroke thickness
• Slant
• Shape
• Size
• Proportions
• Writing styles

This makes MNIST a useful starting point for understanding image classification.

My goal was therefore not simply:

“Can I classify MNIST?”

It was:

“Can I design a neural network pipeline that learns meaningful patterns while controlling overfitting and evaluating its mistakes?”

That difference changed how I approached the project.

2. Preparing the Data

Before training the neural network, I needed to transform the raw dataset into a format that the model could work with efficiently.

I used three main preprocessing steps:

1. Normalize pixel values
2. Flatten the images
3. Convert labels into one-hot encoded vectors

2.1 Normalizing Pixel Values

The original MNIST pixel values range from:

0 → 255

Neural networks generally train more efficiently when numerical inputs are kept on a smaller and consistent scale.

I therefore divided every pixel by 255.

x_train = x_train.astype("float32") / 255.0
x_test = x_test.astype("float32") / 255.0

After normalization, the pixel values fall into:

0 → 1

For example:

0     → 0.0
128   → 0.502
255   → 1.0

This makes the numerical scale much easier for the network to work with.

3. Flattening the Images

Each MNIST image has dimensions:

28 × 28

That means every image contains:

28 × 28 = 784

pixel values.

Because I was building a fully connected neural network rather than a CNN, I converted every 2D image into a one-dimensional vector.

x_train_flat = x_train.reshape(-1, 28 * 28)
x_test_flat = x_test.reshape(-1, 28 * 28)

print(x_train_flat.shape)
print(x_test_flat.shape)

The resulting shapes are:

(60000, 784)
(10000, 784)

So instead of seeing an image as:

28 × 28

the Dense network receives:

784 features

This approach does have a limitation.

Flattening removes the explicit spatial relationship between neighboring pixels. The network no longer knows that two pixels were originally next to each other in the image.

That is one reason CNNs are generally better suited for image problems.

However, I deliberately used a DNN here because I wanted to understand the fundamental mechanics of dense neural networks before moving to convolutional architectures.

4. Converting Labels into One-Hot Encoding

The original labels are integers:

0, 1, 2, ..., 9

For a 10-class classification problem using categorical crossentropy, I converted those labels into one-hot encoded vectors.

from tensorflow.keras.utils import to_categorical

y_train_cat = to_categorical(y_train, num_classes=10)
y_test_cat = to_categorical(y_test, num_classes=10)

For example, the digit 3 becomes:

[0, 0, 0, 1, 0, 0, 0, 0, 0, 0]

The position containing 1 represents the correct class.

This format works naturally with the model’s final Softmax layer and categorical crossentropy loss.

5. Designing the Neural Network

For the model architecture, I used a progressive structure:

784
 ↓
512
 ↓
256
 ↓
128
 ↓
10

The network becomes narrower as it moves deeper.

My complete architecture was:

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

def build_model(input_dim=784, num_classes=10):
    model = keras.Sequential([
        # Hidden Layer 1
        layers.Dense(
            512,
            input_shape=(input_dim,),
            name="dense_1"
        ),
        layers.BatchNormalization(name="batchnorm_1"),
        layers.Activation("relu", name="relu_1"),
        layers.Dropout(0.3, name="dropout_1"),
        # Hidden Layer 2
        layers.Dense(
            256,
            name="dense_2"
        ),
        layers.BatchNormalization(name="batchnorm_2"),
        layers.Activation("relu", name="relu_2"),
        layers.Dropout(0.3, name="dropout_2"),
        # Hidden Layer 3
        layers.Dense(
            128,
            name="dense_3"
        ),
        layers.BatchNormalization(name="batchnorm_3"),
        layers.Activation("relu", name="relu_3"),
        layers.Dropout(0.2, name="dropout_3"),
        # Output Layer
        layers.Dense(
            num_classes,
            activation="softmax",
            name="output"
        )
    ], name="MNIST_DNN")
    return model

I then created the model:

model = build_model()

model.summary()

The model contains approximately 571,018 total parameters, including trainable and non-trainable Batch Normalization parameters.

6. Why 512 → 256 → 128?

The decreasing number of neurons was intentional.

The first hidden layer contains:

512 neurons

The second:

256 neurons

The third:

128 neurons

This creates a funnel-like architecture.

The earlier layers have more capacity to learn different combinations of pixel patterns.

Conceptually, the network can learn representations related to:

• Edges
• Curves
• Stroke intersections
• Horizontal lines
• Diagonal patterns
• Loops

As the information moves deeper into the network, the representation becomes more compact.

The deeper layers can combine simpler patterns into more meaningful representations of the digits.

This gives me a progressively compressed representation:

Raw Pixels
    ↓
Low-Level Patterns
    ↓
Combined Features
    ↓
Compact Representation
    ↓
Digit Probability

7. Why Batch Normalization?

After every Dense layer, I added Batch Normalization.

For example:

layers.Dense(512),
layers.BatchNormalization(),
layers.Activation("relu")

Batch Normalization helps stabilize the values flowing through the network during training.

In my architecture, the sequence is:

Dense
 ↓
Batch Normalization
 ↓
ReLU
 ↓
Dropout

This combination provides both normalization and regularization.

The Batch Normalization layers also introduce non-trainable parameters associated with their running statistics. In my model, these contributed to the reported 1,792 non-trainable parameters.

8. Why ReLU?

For the hidden layers, I used the ReLU activation function:

f(x) = max(0, x)

In Keras:

layers.Activation("relu")

ReLU is widely used in neural networks because it is computationally simple and avoids the saturation behavior associated with sigmoid and tanh for positive inputs.

For my model, ReLU was a natural choice for the hidden layers.

9. Why Dropout?

A neural network with hundreds of thousands of parameters can potentially memorize the training data.

That is where Dropout becomes useful.

I used:

30% Dropout
30% Dropout
20% Dropout

The implementation was:

layers.Dropout(0.3)
layers.Dropout(0.3)
layers.Dropout(0.2)

During training, Dropout randomly deactivates a portion of neurons.

This forces the network to avoid depending too heavily on specific neurons and encourages it to learn more robust representations.

At inference time, Dropout is automatically disabled and the complete network is used.

The decreasing dropout rate was intentional: I applied stronger regularization in the earlier layers and slightly less toward the output.

10. The Softmax Output Layer

The final layer contains 10 neurons:

layers.Dense(
    10,
    activation="softmax"
)

Each neuron represents one digit:

0 → probability
1 → probability
2 → probability
...
9 → probability

For example, the model might produce something conceptually like:

0: 0.001
1: 0.002
2: 0.003
3: 0.991
4: 0.001
...

The probabilities sum to approximately 1.

The class with the highest probability becomes the model’s prediction.

11. Compiling the Model

Once the architecture was defined, I compiled the model using the Adam optimizer and categorical crossentropy loss.

model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.001),
    loss="categorical_crossentropy",
    metrics=["accuracy"]
)

Adam was a practical choice because it adapts the learning rate for individual parameters while incorporating ideas related to momentum and RMSprop.

For this dataset, a learning rate of:

0.001

provided a good starting point.

But I did not want the learning rate to remain fixed throughout the entire training process.

That is where callbacks became important.

12. Making Training Smarter with Callbacks

Instead of simply telling the model:

“Train for 20 epochs no matter what.

I wanted the training process to react to what was happening.

I used two callbacks:

1. Early Stopping
2. ReduceLROnPlateau

12.1 EarlyStopping

My EarlyStopping configuration was:

early_stop = keras.callbacks.EarlyStopping(
    monitor="val_loss",
    patience=5,
    restore_best_weights=True
)

This means the model monitors validation loss.

If the validation loss does not improve for five consecutive epochs, training stops.

The most important part for me was:

restore_best_weights=True

This restores the weights from the epoch where the validation loss was best.

So if training continues to improve for a while and then starts to overfit, I do not have to keep the final, potentially worse set of weights.

13. ReduceLROnPlateau

The second callback dynamically reduces the learning rate when progress slows down.

lr_reduce = keras.callbacks.ReduceLROnPlateau(
    monitor="val_loss",
    factor=0.5,
    patience=3,
    min_lr=1e-6,
    verbose=1
)

The logic is straightforward:

Get Farjana Ferdausi’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

If validation loss stops improving for three epochs:

Learning Rate
     ↓
50%

For example:

0.001
 ↓
0.0005
 ↓
0.00025

until the minimum learning rate is reached.

This allows the optimizer to take smaller steps when the model gets closer to convergence.

14. Training the Model

I trained the network for a maximum of 20 epochs using a batch size of 128.

I also reserved 15% of the training data for validation.

history = model.fit(
    x_train_flat,
    y_train_cat,
    epochs=20,
    batch_size=128,
    validation_split=0.15,
    callbacks=[
        early_stop,
        lr_reduce
    ],
    verbose=1
)

The 15% validation split represents:

60,000 × 0.15 = 9,000 images

This allowed me to monitor how well the model was performing on data that was not directly used for updating the model’s weights.

15. What Happened During Training?

The training history showed a healthy learning pattern.

Training accuracy increased rapidly during the first few epochs.

It reached approximately:

95%

within the early stages and eventually approached:

99%

Validation accuracy followed closely and plateaued around:

98.5%

The training loss decreased substantially, while validation loss also continued to improve.

The relatively small gap between training and validation performance was encouraging.

If training accuracy had continued climbing while validation accuracy remained flat or declined, that would have been a stronger indication of overfitting.

Instead, the curves suggested that the model was learning useful patterns that generalized reasonably well to unseen data.

16. Evaluating the Test Set

After training, I evaluated the model on the completely held-out test set.

test_loss, test_accuracy = model.evaluate(
    x_test_flat,
    y_test_cat,
    verbose=0
)

print(f"Test Loss: {test_loss:.4f}")
print(f"Test Accuracy: {test_accuracy:.4f}")

My final reported performance was approximately:

Test Loss: ~0.06
Test Accuracy: ~98.5%

That means the model correctly classified approximately:

9,850 out of 10,000

test images.

For a fully connected network without convolutional layers, I consider this a strong result.

17. Accuracy Alone Isn’t Enough

One lesson I took from this project is that accuracy is not enough to understand a classifier.

A model can achieve a high overall accuracy while still performing poorly on particular classes.

That is why I also examined the confusion matrix.

A confusion matrix helps answer questions such as:

Which digits does my model confuse with each other?

18. Understanding the Confusion Matrix

The confusion matrix was strongly concentrated along the diagonal, meaning most predictions were correct.

The remaining errors tended to occur between visually similar digits.

Examples included:

4 vs. 9

Some handwritten versions of these digits share similar loop structures.

5 vs. 3

Both can contain curved lower sections and similar upper strokes.

7 vs. 1

A particular handwritten style of 7 can resemble 1.

8 vs. 3

Both can contain continuous curved structures, especially when the 8 is written incompletely.

These mistakes were interesting because they showed that the model was not simply failing randomly.

The difficult examples often contained genuinely ambiguous handwriting characteristics.

19. Looking at Prediction Confidence

I also wanted to understand how confident the model was when making predictions.

Instead of returning only:

Prediction = 3

I extracted the probability distribution.

import numpy as np

def predict_digit(image_array):
    image_norm = image_array.astype("float32") / 255.0
    image_flat = image_norm.reshape(1, 784)
    probabilities = model.predict(
        image_flat,
        verbose=0
    )
    predicted_class = np.argmax(probabilities)
    confidence = np.max(probabilities)
    return predicted_class, confidence

This gives me two useful pieces of information:

Predicted digit
Confidence score

For example:

Prediction: 3
Confidence: 99.4%

High-confidence predictions were generally correct.

Errors tended to appear on more ambiguous handwriting examples — such as digits with unusual slant, broken strokes, or uncommon proportions.

This is particularly important when thinking about production systems.

A real application does not always have to make an automatic decision.

A better design could be:

High confidence
      ↓
Automatic prediction

Low confidence
      ↓
Human review

That is much safer than blindly trusting every prediction.

20. Saving the Trained Model

Once the model was trained, I saved it in Keras format.

model.save("mnist_dnn_model.keras")

This allows me to reuse the trained model later without retraining it from scratch.

For example, I can load it with:

from tensorflow import keras

model = keras.models.load_model(
    "mnist_dnn_model.keras"
)

This is an important step when moving from experimentation toward deployment.

21. Thinking Beyond the Notebook

Achieving 98.5% accuracy inside a notebook is a useful milestone.

But a production system introduces additional challenges.

Real-world handwritten images may not look like MNIST.

They can have:

• Different resolutions
• Background noise
• Uneven lighting
• Different stroke thickness
• Off-center digits
• Different writing instruments
• Unusual proportions

Therefore, a model that performs well on MNIST is not automatically ready for real-world handwriting recognition.

This distinction is important:

Good benchmark performance does not automatically mean production readiness.

22. Improving Generalization with Data Augmentation

For a production-oriented version of this project, I would introduce data augmentation.

Potential transformations include:

Rotation: ±10°
Width shift: ±10%
Height shift: ±10%
Zoom: ±10%
Elastic deformation

The idea is to expose the model to variations that resemble real-world handwriting.

Instead of always showing the network perfectly centered digits, augmentation can teach it to tolerate small changes in:

• Position
• Rotation
• Scale
• Shape

This would be one of the next improvements I would test.

23. The Biggest Limitation: It’s Still a DNN

Although my DNN achieved approximately 98.5% test accuracy, there is an important architectural limitation.

I flattened every image:

28 × 28
   ↓
784

This means the network does not explicitly preserve spatial relationships.

For image classification, spatial structure matters.

A pixel next to another pixel often carries more meaningful information than two pixels located far apart.

This is where Convolutional Neural Networks become particularly useful.

24. My Next Upgrade: CNN

The natural next step would be to replace the flattening-first approach with convolutional layers.

A simplified architecture could look like:

Input Image
     ↓
Conv2D
     ↓
MaxPooling
     ↓
Conv2D
     ↓
MaxPooling
     ↓
Flatten
     ↓
Dense
     ↓
Softmax

A CNN can preserve local spatial patterns while also using weight sharing to reduce the number of parameters.

For MNIST, a simple LeNet-style CNN can typically achieve accuracy above 99%, making it a natural benchmark for the next stage of this project.

25. What This Project Actually Taught Me

The most valuable part of this project was not the final accuracy number.

It was learning how individual engineering decisions affect a machine learning system.

1. Preprocessing matters

Normalizing the input data gives the model a much more manageable numerical range.

x_train = x_train.astype("float32") / 255.0

A seemingly small preprocessing decision can have a significant effect on training behavior.

2. Architecture matters

I did not randomly choose:

512 → 256 → 128

The architecture progressively compresses the learned representation.

3. Regularization matters

My model had more than half a million parameters.

Without appropriate regularization, a model with this capacity could memorize the training data.

Batch Normalization and Dropout helped me control that risk.

4. Training should be adaptive

Instead of blindly running a fixed number of epochs, I used:

EarlyStopping
+
ReduceLROnPlateau

This made the training process responsive to validation performance.

5. Evaluation should go beyond accuracy

A single number such as:

98.5%

does not tell the entire story.

The confusion matrix and confidence scores helped me understand where the model struggled.

26. From MNIST to Real Computer Vision

MNIST is often described as a beginner dataset.

I understand why.

It is small, clean, and easy to experiment with.

But the concepts I practiced here are not limited to handwritten digits.

The same principles appear in larger computer vision systems:

Data Preparation
      ↓
Feature Learning
      ↓
Architecture Design
      ↓
Regularization
      ↓
Optimization
      ↓
Evaluation
      ↓
Deployment

These principles can be applied to problems such as:

• Optical Character Recognition
• Medical image classification
• Industrial defect detection
• Document analysis
• Object recognition
• Automated inspection systems

The dataset may change.

The engineering mindset does not.

27. What I Would Build Next

There are several directions I would take this project.

CNN Architecture

Replace the fully connected architecture with Conv2D and MaxPooling layers.

Data Augmentation

Introduce realistic transformations to improve robustness.

Hyperparameter Tuning

Experiment with:

• Learning rates
• Batch sizes
• Network depth
• Dropout rates
• Number of neurons

Tools such as KerasTuner could help automate this process.

Model Quantization

Explore converting 32-bit model weights to lower-precision representations for more efficient deployment.

Real-World Dataset

Test the model on a more challenging dataset such as EMNIST or SVHN to understand how well the learned approach generalizes beyond MNIST.

These are not simply ways to increase the accuracy score.

They are opportunities to move from:

“I trained a neural network.”

to:

“I engineered a more robust machine learning system.”

28. Final Thoughts

Building this MNIST classifier gave me a much deeper understanding of what happens inside a neural network.

I learned that good model performance is rarely the result of one magical technique.

It comes from a sequence of thoughtful decisions:

Clean Data
   ↓
Good Preprocessing
   ↓
Reasonable Architecture
   ↓
Regularization
   ↓
Adaptive Training
   ↓
Careful Evaluation
   ↓
Production Thinking

My final model achieved approximately 98.5% accuracy on unseen MNIST test data.

But for me, the bigger achievement was understanding the engineering decisions behind that number.

MNIST may be a small dataset, but it provided a practical environment for learning some of the foundations I will need as I continue developing deeper expertise in AI, Machine Learning, and Computer Vision.

The next step is clear:

move from fully connected networks to convolutional architectures, introduce stronger real-world augmentation, and continue pushing this project from a notebook experiment toward a production-ready computer vision system.

📎 Project Resources

Google Colab:
https://colab.research.google.com/drive/1Uyk53Rb4i5C4myr5qPS4S_VHGbc781YZ?usp=sharing

Connect with me on LinkedIn:
https://www.linkedin.com/in/farjana-ferdausi/

If this article was useful to you, feel free to share it with your network. I would also love to hear your thoughts, questions, or suggestions about improving this project.

Thank you for reading.