Image Classification in TensorFlow CIFAR-10 in Python

CIFAR-10 dataset is a collection of images used for object recognition and image classification. CIFAR stands for the Canadian Institute for Advanced Research.

There are 60,000 images with size 32X32 color images which are further divided into 50,000 training images and 10,000 testing images. These images are classified into 10 classes with 6000 images per class.

Our task at hand is to train the ML model using TensorFlow on these images to classify them. This is a classification project.

What is TensorFlow?

Tensorflow is an open-source machine learning library developed by Google. The CIFAR-10 dataset and the ML model is taken from Tensorflow libraries only.

Let’s Code in Python

Note- This project was done in Jupyter Notebook you can use any platform you are comfortable with.

The notebook containing the details to download the data is enclosed.

Firstly import all the required libraries.

import cifar10
import matplotlib.pyplot as plt
import tensorflow as tf

from tensorflow.keras import datasets, layers, models
import numpy as np

Collect the data.

cifar10.data_path = "data/CIFAR-10/"

Now let’s take a look at the class names

class_names = cifar10.load_class_names()

Loading data: data/CIFAR-10/cifar-10-batches-py/batches.meta

Now we will load the training and testing images data, classes and labels.

images_train, cls_train, labels_train = cifar10.load_training_data()
images_test, cls_test, labels_test = cifar10.load_test_data()

(will be loaded in batches)

The images are in the form of a 4-dimensional array.


array([[[[0.23137255, 0.24313725, 0.24705882],
         [0.16862745, 0.18039216, 0.17647059],
         [0.19607843, 0.18823529, 0.16862745],
         [0.61960784, 0.51764706, 0.42352941],
         [0.59607843, 0.49019608, 0.4       ],
         [0.58039216, 0.48627451, 0.40392157]],

        [[0.0627451 , 0.07843137, 0.07843137],
         [0.        , 0.        , 0.        ],
         [0.07058824, 0.03137255, 0.        ],
         [0.48235294, 0.34509804, 0.21568627],
         [0.46666667, 0.3254902 , 0.19607843],
         [0.47843137, 0.34117647, 0.22352941]],

        [[0.09803922, 0.09411765, 0.08235294],
         [0.0627451 , 0.02745098, 0.        ],
         [0.19215686, 0.10588235, 0.03137255],

These 4 dimensions represent an array of 50000 images, which are again 32X32, each image is an array of pixels, and every pixel is inturn an array of 3, RGB values.

This is shown in this command

(50000, 32, 32, 3)

Similarly for testing data.

If you wish to see the image, you can do it in this way.



Let’s get our Neural Network ready. We will add 5 layers to the network.

The first layer takes the input data. We will put 2 convolution layers and 2 max-pooling layers, alternatively.

model = models.Sequential()
model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 3)))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))


model.add(layers.Dense(64, activation='relu'))
model.add(layers.Dense(10, activation='softmax'))
model.compile(optimizer='adam',loss='sparse_categorical_crossentropy', metrics=['accuracy'])

The metric used to judge the model is accuracy.

Now let’s train the model, cls_train, epochs=10 ,validation_data=(images_test, cls_test))

The next step is to test our model and get the predictions.

predictions = model.predict(images_test)
score = model.evaluate(images_test, cls_test)

Our Image Classification model is ready.

Leave a Reply

Your email address will not be published. Required fields are marked *