How to fine-tune the pre-trained VGG model in TensorFlow Keras

In this tutorial, we will demonstrate how to use the previously fine-tune trained VGG16 model in TensorFlow Keras to classify our own image.

VGG16 won the 2014 ImageNet competition this is basically computation where there are 1000 of images belonging to 1000 different categories.VGG model weights are freely available and can be loaded and used in your own models and applications. This allowed other researchers and developers to use a state-of-the-art image classification model in their own work and programs.

Download Data

Before you start, you’ll need a set of images to teach the network about the new classes you want to recognize. Google has created an archive of creative-commons’ licensed flower photos to use initially.

from __future__ import absolute_import, division, print_function
from tqdm import tqdm
from numpy.random import randn

import pathlib
import random
import matplotlib.pyplot as plt

import tensorflow as tf
import numpy as np

from matplotlib.image import imread
from keras.preprocessing import image

tf.enable_eager_execution()

AUTOTUNE = tf.data.experimental.AUTOTUNE

data_dir = tf.keras.utils.get_file('flower_photos','https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz', untar=True)
data_dir = pathlib.Path(data_dir)

label_names={'daisy': 0, 'dandelion': 1, 'roses': 2, 'sunflowers': 3, 'tulips': 4}
label_key=['daisy','dandelion','roses','sunflowers','tulips']

Load images with tf.data

In this tutorial, we will use tf.data API to load data into the model. The easiest way to build a tf.data.Dataset is using the from_tensor_slices method.

all_images = list(data_dir.glob('*/*'))
all_images = [str(path) for path in all_images]
random.shuffle(all_images)

all_labels=[label_names[pathlib.Path(path).parent.name] for path in all_images]

data_size=len(all_images)

train_test_split=(int)(data_size*0.2)

x_train=all_images[train_test_split:]
x_test=all_images[:train_test_split]

y_train=all_labels[train_test_split:]
y_test=all_labels[:train_test_split]

IMG_SIZE=160

BATCH_SIZE = 32

def _parse_data(x,y):
  image = tf.read_file(x)
  image = tf.image.decode_jpeg(image, channels=3)
  image = tf.cast(image, tf.float32)
  image = (image/127.5) - 1
  image = tf.image.resize(image, (IMG_SIZE, IMG_SIZE))
 
  return image,y

def _input_fn(x,y):
  ds=tf.data.Dataset.from_tensor_slices((x,y))
  ds=ds.map(_parse_data)
  ds=ds.shuffle(buffer_size=data_size)
  
  
  ds = ds.repeat()
  
  ds = ds.batch(BATCH_SIZE)
  
  ds = ds.prefetch(buffer_size=AUTOTUNE)
  
  return ds
  
train_ds=_input_fn(x_train,y_train)
validation_ds=_input_fn(x_test,y_test)

Create the base model from VGG16-trained convnets.

We will create a base model from the VGG16 model. This is pre-trained on the ImageNet dataset, a large dataset of 1.4M images and 1000 classes of web images.

The very last classification layer is not very useful. Instead, we will follow the common practice to instead depend on the very last layer before the flatten operation. This layer is called the “bottleneck layer”. The bottleneck features retain many generalities as compared to the final/top layer.

First, instantiate a VGG16 model pre-loaded with weights trained on ImageNet. By specifying the include_top=False argument, you load a network that doesn’t include the classification layers.

IMG_SHAPE = (IMG_SIZE, IMG_SIZE, 3)
VGG16_MODEL=tf.keras.applications.VGG16(input_shape=IMG_SHAPE,
                                               include_top=False,
                                               weights='imagenet')

It’s important to freeze the convolutional base before you compile and train the model. By freezing or setting layer.trainable = False, you prevent the weights in a given layer from being updated during training.

VGG16_MODEL.trainable=False
global_average_layer = tf.keras.layers.GlobalAveragePooling2D()
prediction_layer = tf.keras.layers.Dense(len(label_names),activation='softmax')

Apply a tf.keras.layers.Dense layer to convert these features into a single prediction per image.

Now stack the feature extractor, and these two layers using a tf.keras.Sequential model.

model = tf.keras.Sequential([
  VGG16_MODEL,
  global_average_layer,
  prediction_layer
])

Compile the model

You need to compile the model before training it. Since there are five classes, use a sparse_categorical_crossentropy.

model.compile(optimizer=tf.train.AdamOptimizer(), 
              loss=tf.keras.losses.sparse_categorical_crossentropy,
              metrics=["accuracy"])

Train the model

Here, the fit method uses the steps_per_epoch argument—this is the number of training steps the model runs before it moves to the next epoch.

history = model.fit(train_ds,
                    epochs=100, 
                    steps_per_epoch=2,
                    validation_steps=2,
                    validation_data=validation_ds)

Evaluate Model

The tf.keras.Model.evaluate methods use NumPy data and a tf.data.Dataset.To evaluate the inference-mode loss and metrics for the data provided.

validation_steps = 20

loss0,accuracy0 = model.evaluate(validation_ds, steps = validation_steps)

print("loss: {:.2f}".format(loss0))
print("accuracy: {:.2f}".format(accuracy0))

Learning curves

Let’s take a look at the learning curves of the training and validation accuracy/loss when using the VGG16 base model.

plt.plot(history.history['acc'])
plt.plot(history.history['val_acc'])
plt.title('model accuracy')
plt.ylabel('accuracy')
plt.xlabel('epoch')
plt.legend(['train', 'test'], loc='upper left')
plt.show()

# summarize history for loss
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['train', 'test'], loc='upper left')
plt.show()