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

Selectcaret down icon
Select Area of interestcaret down icon
Select Work Experiencecaret down icon
By clicking 'Submit' you Agree to  
UpGrad's Terms & Conditions

Our Popular Machine Learning Course

Explore Free Courses

Suggested Blogs

15 Interesting MATLAB Project Ideas & Topics For Beginners [2024]
Diving into the world of engineering and data science, I’ve discovered the potential of MATLAB as an indispensable tool. It has accelerated my c
Read More

by Pavan Vadapalli

09 Jul 2024

5 Types of Research Design: Elements and Characteristics
The reliability and quality of your research depend upon several factors such as determination of target audience, the survey of a sample population,
Read More

by Pavan Vadapalli

07 Jul 2024

Biological Neural Network: Importance, Components & Comparison
Humans have made several attempts to mimic the biological systems, and one of them is artificial neural networks inspired by the biological neural net
Read More

by Pavan Vadapalli

04 Jul 2024

Production System in Artificial Intelligence and its Characteristics
The AI market has witnessed rapid growth on the international level, and it is predicted to show a CAGR of 37.3% from 2023 to 2030. The production sys
Read More

by Pavan Vadapalli

03 Jul 2024

AI vs Human Intelligence: Difference Between AI & Human Intelligence
In this article, you will learn about AI vs Human Intelligence, Difference Between AI & Human Intelligence. Definition of AI & Human Intelli
Read More

by Pavan Vadapalli

01 Jul 2024

Career Opportunities in Artificial Intelligence: List of Various Job Roles
Artificial Intelligence or AI career opportunities have escalated recently due to its surging demands in industries. The hype that AI will create tons
Read More

by Pavan Vadapalli

26 Jun 2024

Gini Index for Decision Trees: Mechanism, Perfect & Imperfect Split With Examples
As you start learning about supervised learning, it’s important to get acquainted with the concept of decision trees. Decision trees are akin to
Read More

by MK Gurucharan

24 Jun 2024

Random Forest Vs Decision Tree: Difference Between Random Forest and Decision Tree
Recent advancements have paved the growth of multiple algorithms. These new and blazing algorithms have set the data on fire. They help in handling da
Read More

by Pavan Vadapalli

24 Jun 2024

Basic CNN Architecture: Explaining 5 Layers of Convolutional Neural Network
Introduction In the last few years of the IT industry, there has been a huge demand for once particular skill set known as Deep Learning. Deep Learni
Read More

by MK Gurucharan

21 Jun 2024

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