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

Artificial Intelligence course fees
Artificial intelligence (AI) was one of the most used words in 2023, which emphasizes how important and widespread this technology has become. If you
Read More

by venkatesh Rajanala

29 Feb 2024

Artificial Intelligence in Banking 2024: Examples & Challenges
Introduction Millennials and their changing preferences have led to a wide-scale disruption of daily processes in many industries and a simultaneous g
Read More

by Pavan Vadapalli

27 Feb 2024

Top 9 Python Libraries for Machine Learning in 2024
Machine learning is the most algorithm-intense field in computer science. Gone are those days when people had to code all algorithms for machine learn
Read More

by upGrad

19 Feb 2024

Top 15 IoT Interview Questions & Answers 2024 – 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

19 Feb 2024

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 Feb 2024

Artificial Intelligence Salary in India [For Beginners & Experienced] in 2024
Artificial Intelligence (AI) has been one of the hottest buzzwords in the tech sphere for quite some time now. As Data Science is advancing, both AI a
Read More

by upGrad

18 Feb 2024

24 Exciting IoT Project Ideas & Topics For Beginners 2024 [Latest]
Summary: In this article, you will learn the 24 Exciting IoT Project Ideas & Topics. Take a glimpse at the project ideas listed below. Smart Agr
Read More

by Kechit Goyal

18 Feb 2024

Natural Language Processing (NLP) Projects & Topics For Beginners [2023]
What are Natural Language Processing Projects? NLP project ideas advanced encompass various applications and research areas that leverage computation
Read More

by Pavan Vadapalli

17 Feb 2024

45+ Interesting Machine Learning Project Ideas For Beginners [2024]
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

16 Feb 2024

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