Transfer Learning in Keras (Image Recognition)


Transfer Learning in AI is a method where a model is developed for a specific task, which is used as the initial steps for another model for other tasks.

Deep Convolutional Neural Networks in deep learning take an hour or day to train the mode if the dataset we are playing is vast.


The approach is we reuse the weights of the pre-trained model, which was trained for some standard Computer Vision datasets such as Image classification (Image Net).

Extensive deep Convolutional networks for large-scale image classification are available in Keras, which we can directly import and can be used with their pre-trained weights.

Various pre-trained model is available in Keras are:

Let’s now understand how to use VGG16 pre-trained on 10,000 categories(Image Net) for the Distracted driver Detection dataset.

Dataset link:

Dataset has ten categories to classify, but VGG16 was trained for 10,000 categories, so to apply VGG16 to the Distracted Driver dataset, Fully connected layers need some changes. At the same time, the Convolutional network of VGG16 is used as such with their pre-trained weights.


Importing the required library

from tensorflow.keras.layers import Input, Lambda, Dense, Flatten
from tensorflow.keras.models import Model, Sequential
from tensorflow.keras.applications.vgg16 import VGG16 ,preprocess_input     #importing VGG16 model
from tensorflow.keras.preprocessing import image
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import CSVLogger
import numpy as np 
import os
import matplotlib.pyplot as plt
import tensorflow as tf

The code below is to Shift the computation to GPU’s

# --------------------------------------#
#shifting computation to GPU
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    gpu = gpus[0]
    tf.config.experimental.set_memory_growth(gpu, True)
    tf.config.set_visible_devices(gpu, 'GPU')


The code of the main model is given below:

# Path to my train and val dataset
base = os.path.join(os.getcwd(),'imgs')

# include_top = False means that we doesnt include fully connected top layer we will add them accordingly
vgg16 = VGG16(include_top = False, input_shape = (224,224,3), weights = 'imagenet')

# training of all the convolution is set to false
for layer in vgg16.layers:
    layer.trainable = False

x = Flatten()(vgg16.output)
predictions = Dense(units = 10, activation = 'sigmoid')(x)

model = Model(inputs = vgg16.input, outputs = predictions)
model.compile(optimizer = Adam(learning_rate = 0.01 ), loss = 'binary_crossentropy',metrics = ['accuracy'] )

To  load  the dataset batch-wise, we use an in-built Image Data generator:      

train_gen = ImageDataGenerator(rescale = 1./255)
val_gen = ImageDataGenerator(rescale = 1./255)

train_data = train_gen.flow_from_directory(os.path.join(base,'train'), target_size = (244,244),batch_size = 16, class_mode = 'categorical')
val_data = val_gen.flow_from_directory(os.path.join(base,'val'), target_size = (244,244),batch_size = 16, class_mode = 'categorical'),validation_data = val_data,epochs = 30,callbacks = [CSVLogger("distracted_driver.csv",append = True)])

Model summary

Accuracy Trend:

Loss Trend:


Please enter your comment!
Please enter your name here

This site uses Akismet to reduce spam. Learn how your comment data is processed.