MNIST Digit Classification Using TensorFlow in Python

The MNIST dataset is one of the most widely used datasets in machine learning and deep learning. It contains thousands of handwritten digit images and is commonly used as a benchmark for testing image classification models. Because of its simplicity and well-structured format, MNIST is often the starting point for beginners learning neural networks and computer vision.

In this tutorial, we will perform MNIST digit classification using TensorFlow in Python. You’ll learn how to load the MNIST dataset, build a simple neural network model, train it using TensorFlow, and evaluate its performance on handwritten digit recognition.

MNIST Digit Classification Using TensorFlow

This type of classification problem involves building a neural network model that can recognize handwritten digits from images. TensorFlow provides built-in tools that make it easy to load datasets, design neural network architectures, and train models efficiently.

The classification process typically begins by loading the MNIST dataset using tf.keras.datasets.mnist. After loading the dataset, it is divided into a training set and a test set. The training data is used to teach the model patterns in handwritten digits, while the test data evaluates how well the model performs on unseen images.

Before training, the image pixel values are usually normalized between 0 and 1. This preprocessing step helps the neural network learn faster and improves training stability.

A common approach for MNIST digit classification is to build a simple neural network using

TensorFlow’s Keras API. The model typically includes:

1. A Flatten layer to convert the 28 × 28 image into a vector

2. One or more Dense layers with ReLU activation

3. An output layer with Softmax activation to predict probabilities for the digits 0–9

Once the architecture is defined, the model is compiled using an optimizer such as Adam and a loss function like sparse categorical crossentropy, which is suitable for multi-class classification tasks.

The model is then trained for several epochs, allowing it to adjust its weights and improve its predictions. After training, the model is evaluated on the test dataset to measure its accuracy. The trained model can then classify new handwritten digit images.

The MNIST dataset (Modified National Institute of Standards and Technology) is one of the most widely used datasets in machine learning. It contains 70,000 grayscale images of handwritten digits ranging from 0 to 9.

Each image is 28 × 28 pixels, resulting in 784 total pixels per image. Pixel values range from 0 to 255, where 0 represents black and 255 represents white.

The dataset is divided into two parts:

Training Set with 60,000 images used to train machine learning models.

Test Set with 10,000 images used to evaluate model performance.

MNIST is commonly used for learning and benchmarking image classification algorithms because it is small, well-structured, and easily accessible through machine learning libraries such as TensorFlow.

Implementing Deep Learning Neural Networks over MNIST Dataset

Now that we understand the MNIST dataset and its role in handwritten digit recognition, the next step is to build a neural network model to perform MNIST digit classification using TensorFlow. Deep learning models are particularly effective for image classification tasks because they can automatically learn patterns and features from image data.

Using TensorFlow’s Keras API, we can create a neural network with multiple layers, train it on the MNIST dataset, and evaluate its performance on unseen data. The implementation process typically involves loading the dataset, preprocessing the images, defining the neural network architecture, training the model, and finally evaluating its accuracy.

In the following steps, we will walk through the complete process of implementing a deep learning neural network to classify handwritten digits from the MNIST dataset.

1. Loading the MNIST Dataset with TensorFlow

TensorFlow provides direct access to the MNIST dataset through the tf.keras.datasets module. This allows the dataset to be loaded quickly without manually downloading or preparing the files. After loading the data, the next step is to normalize the pixel values so the neural network can train more efficiently.

import tensorflow as tf

# Load the MNIST dataset
mnist = tf.keras.datasets.mnist

# Split into training and testing datasets
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# Normalize pixel values
x_train, x_test = x_train / 255.0, x_test / 255.0

The load_data() function loads the dataset and separates it into training and testing sets. The image data is stored in x_train and x_test, while the corresponding labels are stored in y_train and y_test. Normalizing the pixel values by dividing them by 255.0 scales the data to a range between 0 and 1, which helps improve the stability and speed of the training process.

Before training the neural network, it is useful to visualize some samples from the dataset. This provides a quick look at the handwritten digits the model will learn to recognize.

import matplotlib.pyplot as plt

for i in range(10):    
	plt.subplot(2, 5, i + 1)    
	plt.imshow(x_train[i], cmap='gray')    
	plt.title(f"Label: {y_train[i]}")    
	plt.axis('off')plt.show()

Running this code displays several handwritten digit images along with their labels, giving a quick visual confirmation that the dataset has been loaded and prepared correctly.

2. Building a Simple Neural Network with TensorFlow

After loading and preparing the dataset, the next step is to build a neural network that can classify handwritten digits. TensorFlow’s Keras API makes it straightforward to define neural network architectures using a sequential model.

In this example, the network consists of a flatten layer that converts the image data into a one-dimensional vector, followed by two dense hidden layers that learn patterns in the data. The final layer uses a softmax activation function to predict the probability of each digit class.

# Build the model
model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)), # Flatten the input
    tf.keras.layers.Dense(128, activation='relu'), # First hidden layer
    tf.keras.layers.Dense(64, activation='relu'), # Second hidden layer
    tf.keras.layers.Dense(10, activation='softmax') # Output layer
])

# Compile the model
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

The Flatten layer converts each image into a one-dimensional array so it can be processed by the dense layers. The hidden layers use the ReLU activation function to learn complex patterns in the input data. The output layer contains ten neurons, each representing a possible digit from 0 to 9, and the softmax function converts the output into probability values.

The model is then compiled using the Adam optimizer and sparse categorical crossentropy as the loss function, which is suitable for multi-class classification problems such as handwritten digit recognition.

4. Training  Neural Network Model

Once the neural network architecture is defined and compiled, the next step is to train the model using the training dataset. During training, the model processes the input images and adjusts its internal weights to minimize the loss function and improve prediction accuracy.

Training is performed using the fit() function, which feeds the training data into the neural network for a specified number of epochs.

# Train the model
history = model.fit(x_train, y_train, epochs=5)

Epoch 1/5
1875/1875 ━━━━━━━━━━━━━━━━━━━━ 10s 4ms/step - accuracy: 0.8792 - 
.
.
.
1875/1875 ━━━━━━━━━━━━━━━━━━━━ 6s 3ms/step - accuracy: 0.9870 - loss: 0.0392

Each epoch represents one complete pass through the entire training dataset. As training progresses, the model gradually improves its ability to recognize handwritten digits by learning patterns in the data.

As the epochs progress, the accuracy increases while the loss decreases, indicating that the neural network is learning to classify the handwritten digits more effectively.

5. Evaluating the Model Performance

After training the neural network, the next step is to evaluate how well the model performs on unseen data. This is done using the testing dataset, which was not used during training. Evaluating the model on this dataset provides an estimate of how accurately the neural network can classify new handwritten digits.

The evaluate() function calculates the model’s loss and accuracy on the test data.

# Evaluate the model
test_loss, test_acc = model.evaluate(x_test, y_test)
print('\nTest accuracy:', test_acc)

The evaluate() method returns two values. The first value represents the loss, which indicates how well the model fits the test data. The second value represents the accuracy, showing the percentage of correctly classified images.

After running this code, TensorFlow will display the model’s performance on the test dataset along with the final test accuracy. A well-trained MNIST classification model typically achieves accuracy about 97%, demonstrating the effectiveness of neural networks for handwritten digit recognition.

313/313 ━━━━━━━━━━━━━━━━━━━━ 1s 1ms/step - accuracy: 0.9697 - loss: 0.1007

Test accuracy: 0.9697

6. Making Predictions

After training and evaluating the neural network, the trained model can be used to make predictions on new data. In this step, the model analyzes the input images and predicts the most likely digit for each image.

The predict() function generates probability scores for each digit class. The digit with the highest probability is selected as the final prediction.

# Make predictions
predictions = model.predict(x_test)

# Print some predictions
for i in range(5):
    print('Predicted value: ', predictions[i].argmax())

Output:
Predicted value: 7
Predicted value: 2
Predicted value: 1
Predicted value: 0
Predicted value: 4

The model analyzes each test image and returns the predicted digit. The argmax() function selects the index with the highest probability, which corresponds to the predicted digit class.

This demonstrates how a trained neural network can successfully recognize handwritten digits and make predictions on previously unseen images.

7. Visualizing Model Performance

After training and evaluating the neural network, it is helpful to visualize how the model performed during training. Plotting the accuracy and loss values across epochs provides insight into how well the model learned from the training data.

Training accuracy shows how well the model performed on the training dataset, while training loss indicates how the error changed as the model improved. Comparing these values helps determine whether the model is learning effectively.

# Plot training & validation accuracy values
plt.figure(figsize=(12, 5))

plt.subplot(1, 2, 1)
plt.plot(history.history['accuracy'])
plt.title('Model accuracy')
plt.ylabel('Accuracy')
plt.xlabel('Epoch')
plt.legend(['Train'], loc='upper left')
plt.axhline(y=test_acc, color='r', linestyle='--', label=f'Test Accuracy: {test_acc:.4f}')
plt.legend()

# Plot training & validation loss values
plt.subplot(1, 2, 2)
plt.plot(history.history['loss'])
plt.title('Model loss')
plt.ylabel('Loss')
plt.xlabel('Epoch')
plt.legend(['Train'], loc='upper left')
plt.axhline(y=test_loss, color='r', linestyle='--', label=f'Test Loss: {test_loss:.4f}')
plt.legend()

plt.tight_layout()
plt.show()

The accuracy graph shows how the model’s performance improved over multiple training epochs. The loss graph illustrates how the model’s prediction error decreased during training. The dashed red line represents the final performance of the model on the test dataset.

Visualizing these metrics helps confirm that the neural network trained properly and provides insight into the overall learning behavior of the model.

Conclusion

In this tutorial, we learned how to perform MNIST digit classification using TensorFlow in Python. We started by loading and visualizing the dataset, then built a simple neural network using TensorFlow’s Keras API. The model was trained on handwritten digit images and evaluated to measure its accuracy on test data.

This example demonstrates how deep learning models can effectively solve image classification problems. While the neural network used here is relatively simple, the same approach can be extended to more advanced architectures and larger datasets for real-world machine learning applications.