Homebreadcumb forward arrow iconBlogbreadcumb forward arrow iconArtificial Intelligencebreadcumb forward arrow iconTensorflow 2.0 Image Classification: Install, Load Data, Building & Training the Model

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

Last updated:
21st Apr, 2020
Read Time
5 Mins
share image icon
In this article
Chevron in toc
View All
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.

Top Machine Learning and AI Courses Online

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. 

Ads of upGrad blog

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 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 objects:

train_data =

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


val_data =

 (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 =

   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 = (




val_data = (





Trending Machine Learning Skills

Enrol for the Machine Learning Course from the World’s top Universities. Earn Masters, Executive PGP, or Advanced Certificate Programs to fast-track your career.

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(





# 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

Learn data science courses from the World’s top Universities. Earn Executive PG Programs, Advanced Certificate Programs, or Masters Programs to fast-track your career.

Training the Model

After we’ve built the model, we can teach it. The tf.keras API of TensorFlow 2.0 supports the API, so you have to use the 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.

Popular AI and ML Blogs & Free Courses

Ads of upGrad blog

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.


Kechit Goyal

Blog Author
Experienced Developer, Team Player and a Leader with a demonstrated history of working in startups. Strong engineering professional with a Bachelor of Technology (BTech) focused in Computer Science from Indian Institute of Technology, Delhi.
Get Free Consultation

Select Coursecaret down icon
Selectcaret down icon
By clicking 'Submit' you Agree to  
UpGrad's Terms & Conditions

Our Popular Machine Learning Course

Explore Free Courses

Suggested Blogs

Top 8 Exciting AWS Projects & Ideas For Beginners [2023]
AWS Projects & Topics Looking for AWS project ideas? Then you’ve come to the right place because, in this article, we’ve shared multiple AWS proj
Read More

by Pavan Vadapalli

19 Sep 2023

Data Preprocessing in Machine Learning: 7 Easy Steps To Follow
Summary: In this article, you will learn about data preprocessing in Machine Learning: 7 easy steps to follow. Acquire the dataset Import all the cr
Read More

by Kechit Goyal

18 Sep 2023

Top 15 IoT Interview Questions & Answers 2023 – For Beginners & Experienced
These days, the minute you indulge in any technology-oriented discussion, interview questions on cloud computing come up in some form or the other. Th
Read More

by Kechit Goyal

15 Sep 2023

45+ Interesting Machine Learning Project Ideas For Beginners [2023]
Summary: In this Article, you will learn Stock Prices Predictor Sports Predictor Develop A Sentiment Analyzer Enhance Healthcare Prepare ML Algorith
Read More

by Jaideep Khare

14 Sep 2023

25 Machine Learning Interview Questions & Answers – Linear Regression
Introduction Machine Learning Interviews can vary according to the types or categories, for instance, a few recruiters ask many Linear Regression int
Read More

by Thulasiram Gunipati

10 Sep 2023

13 Interesting Neural Network Project Ideas & Topics for Beginners [2023]
The topic of neural networks has captivated the world of artificial intelligence and machine learning with its ability to mimic the human brain’
Read More

by Pavan Vadapalli

07 Sep 2023

14 Raspberry Pi Project Ideas & Topics For Beginners in 2023
What is a Raspberry Pi? The Raspberry Pi is a low-cost computer about the size of a credit card that can be connected to a display or TV and controll
Read More

by Kechit Goyal

06 Sep 2023

AWS Salary in India in 2023 [For Freshers & Experienced]
Summary: In this article, you will learn about AWS Salary in India For Freshers & Experienced. AWS Salary in India INR 6,07,000 per annum AW
Read More

by Pavan Vadapalli

04 Sep 2023

9 Interesting Linear Regression Project Ideas & Topics For Beginners [2023]
Linear regression is a popular topic in machine learning. It’s a supervised learning algorithm and finds applications in many sectors. If you’re learn
Read More

by Pavan Vadapalli

02 Sep 2023

Schedule 1:1 free counsellingTalk to Career Expert
footer sticky close icon