Tensorflow 2.0 Image Classification: Install, Load Data, Building & Training the Model

Image classification is a category of pattern recognition. It classifies images according to the relationship between the neighboring pixels. In other words, it uses contextual information to organize images and is quite popular among different technologies. It’s a prominent topic in Deep Learning, and if you’re learning about it, you’d surely enjoy this article.

Here, we’ll perform the TensorFlow image classification. We’ll build a model, train it, and then enhance its accuracy to classify images of cacti. TensorFlow is an open-source machine learning platform and a product of Google. 

Let’s get started. 

Install TensorFlow 2.0

First, you’ll need to install TensorFlow on Google Colab. You can install it through pip:

!pip install tensorflow-gpu==2.0.0-alpha0

Then we’ll verify the installation:

import tensorflow as tf


# Output: 2.0.0-alpha0


Learn: Most Popular 5 TensorFlow Projects for Beginners

Load Data

After the verification, we can load the data by using the tf.data.dataset. We’ll build a classifier that determines whether an image contains a cactus or not. The cactus has to be columnar.We can use the Cactus Aerial Photos dataset for this purpose. Now, we’ll load the file paths along with their labels:

train_csv = pd.read_csv(‘data/train.csv’)

# Prepend image filenames in train/ with relative path

filenames = [‘train/’ + fname for fname in train_csv[‘id’].tolist()]

labels = train_csv[‘has_cactus’].tolist()

train_filenames, val_filenames, train_labels, val_labels =





Once we have the labels and filenames, we are ready to create the tf.data.Dataset objects:

train_data = tf.data.Dataset.from_tensor_slices(

 (tf.constant(train_filenames), tf.constant(train_labels))


val_data = tf.data.Dataset.from_tensor_slices(

 (tf.constant(val_filenames), tf.constant(val_labels))



At the moment, our dataset doesn’t have the actual images. It only has their filenames. We’ll need a function to load the necessary images and process them so we can perform TensorFlow image recognition on them.

IMAGE_SIZE = 96 # Minimum image size for use with MobileNetV2


# Function to load and preprocess each image

def _parse_fn(filename, label):

   img = tf.io.read_file(img)

   img = tf.image.decode_jpeg(img)

   img = (tf.cast(img, tf.float32)/127.5) – 1

   img = tf.image.resize(img, (IMAGE_SIZE, IMAGE_SIZE))

   return img, label

# Run _parse_fn over each example in train and val datasets

# Also shuffle and create batches

train_data = (train_data.map(_parse_fn)




val_data = (val_data.map(_parse_fn)





Building a Model

In this TensorFlow image classification example, we’ll create a transfer learning model. These models are fast as they can use existing image classification models that have undergone training before. They only have to retrain the upper layer of their network as this layer specifies the class of the required image. 

We’ll use the Keras API of TensorFlow 2.0 to create our image classification model. And for the transfer learning purposes, we’ll use the MobileNetV2 to be the attribute detector. It’s the second version of MobileNet and is a product of Google. It’s lighter in weight than other models such as Inception and ResNet and can run on mobile devices. We’ll load this model on ImageNet, freeze the weights, add a classification head and run it without its top layer. 


# Pre-trained model with MobileNetV2

base_model = tf.keras.applications.MobileNetV2(





Get our free ebook!
A Beginners Guide to Fundamentals of
Natural Language Processing
Download Now

# Freeze the pre-trained model weights

base_model.trainable = False

# Trainable classification head

maxpool_layer = tf.keras.layers.GlobalMaxPooling2D()

prediction_layer = tf.keras.layers.Dense(1, activation=’sigmoid’)

# Layer classification head with feature detector

model = tf.keras.Sequential([





learning_rate = 0.0001

# Compile the model






You should use TensorFlow optimizers if you are going to train tf.keras models. The optimizers in tf.keras.optimizers and tf.train APIs are together in TensorFlow 2.0’s tf.keras.optimizers. In TensorFlow 2.0, many of the original optimizers of tf.keras have received upgrades and replacements for better performance. They enable us to apply optimizers without compromising with the performance and save time as well. 

Read: TensorFlow Object Detection Tutorial For Beginners

Training the Model

After we’ve built the model, we can teach it. The tf.keras API of TensorFlow 2.0 supports the tf.data API, so you have to use the tf.data.Dataset objects for this purpose. It can perform the training efficiently, and we wouldn’t have to make any compromises with the performance. 

num_epochs = 30

steps_per_epoch = round(num_train)//BATCH_SIZE

val_steps = 20



         steps_per_epoch = steps_per_epoch,




After 30 epochs, the model’s accuracy increases substantially, but we can improve it further. Remember, we mentioned freezing the weights during transfer learning? Well, now that we have trained the classification head, we can unfreeze those layers and fine-tune our dataset further:

# Unfreeze all layers of MobileNetV2

base_model.trainable = True

# Refreeze layers until the layers we want to fine-tune

for layer in base_model.layers[:100]:

 layer.trainable =  False

# Use a lower learning rate

lr_finetune = learning_rate / 10

# Recompile the model


             optimizer = tf.keras.optimizers.Adam(lr=lr_finetune),


# Increase training epochs for fine-tuning

fine_tune_epochs = 30

total_epochs =  num_epochs + fine_tune_epochs

# Fine-tune model

# Note: Set initial_epoch to begin training after epoch 30 since we

# previously trained for 30 epochs.


         steps_per_epoch = steps_per_epoch,


         initial_epoch = num_epochs,




30 epochs later, the model’s accuracy improves further. WIth more epochs, we saw more improvement in the accuracy of the model. Now, we have a proper TensorFlow image recognition model that can recognize columnar cacti in images with a high accuracy. 

Also Read: Tensorflow Project Ideas for Beginners

Learn More About TensorFlow Image Classification

The highly functional APIs of TensorFlow and its capabilities make it a powerful technology for any programmer to wield. Its high-level APIs also remove its general complexity, making it easier to use. 

Are you interested in learning more about TensorFlow, image classification, and related topics? Then we recommend you to check out IIIT-B & upGrad’s PG Diploma in Machine Learning & AI which is designed for working professionals and offers 450+ hours of rigorous training, 30+ case studies & assignments, IIIT-B Alumni status, 5+ practical hands-on capstone projects & job assistance with top firms.

Machine learning course | Learn Online, IIIT Bangalore‎

Apply Now

Leave a comment

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

Know More
Download EBook
Download EBook
By clicking Download EBook, you agree to our terms and conditions and our privacy policy.
Get our free ebook!
A Beginners Guide to Fundamentals of
Natural Language Processing
Download Now