Fine-tune VGG16 model for image classification in Keras

Keras framework provides us a lot of pre-trained general purpose deep learning models which we can fine-tune as per our requirements. We don't need to build a complex model from scratch. In my last article, we built a CNN model from scratch for image classification. Instead of that, we can just fine-tune an existing, well-trained, well-proven, widely accepted CNN model which will save our a lot of effort, time and money.

VGG16 is a proven proficient algorithm for image classification (1000 classes of images). Keras framework already contain this model. We will import this model and fine-tune it to classify the images of dogs and cats (only 2 classes instead of 1000 classes).

You can download my Jupyter notebook containing below code from here.

Step 1: Import the required libraries

import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

import keras
from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.layers import Dense
from keras.optimizers import Adam

from sklearn.metrics import confusion_matrix, accuracy_score, classification_report

Step 2: Create directory structure to contain images

We will create a directory structure which will contain the images of dogs and cats.

I have created a directory "cats_and_dogs". Under this directory, I have created 3 other directories "test", "train" and "valid". All these 3 directories contain "cat" and "dog" directories. 

1. "cat" and "dog" directories under "test" directory contain 5 images of cats and dogs respectively. Total 10 images for testing.

2. "cat" and "dog" directories under "train" directory contain 20 images of cats and dogs respectively. Total 40 images for training.

3. "cat" and "dog" directories under "valid" directory contain 8 images of cats and dogs respectively. Total 16 images for validation.

Step 3: Data Preparation

train_path = 'C:/cats_and_dogs/train'
valid_path = 'C:/cats_and_dogs/valid'
test_path = 'C:/cats_and_dogs/test'

train_batches = ImageDataGenerator().flow_from_directory(train_path, target_size=(224,224), classes=['dog','cat'], batch_size=10)

valid_batches = ImageDataGenerator().flow_from_directory(valid_path, target_size=(224,224), classes=['dog','cat'], batch_size=4)

test_batches = ImageDataGenerator().flow_from_directory(test_path, target_size=(224,224), classes=['dog','cat'], batch_size=10)

Output:
Found 40 images belonging to 2 classes. Found 16 images belonging to 2 classes. Found 10 images belonging to 2 classes.

In the above code, we are generating the images of 224x224 pixels and categorizing these images into cat and dog classes. It is clear from the output that we have 40 images for training, 16 images for validation and 10 images for testing as mentioned in step 2.

Step 4: Print the images

Lets output some of the images which we have prepared in step 3. Following is the standard code to print the images (copied from Keras documentation)

def plots(ims, figsize=(12,6), rows=1, interp=False, titles=None):
    if type(ims[0]) is np.ndarray:
        ims = np.array(ims).astype(np.uint8)
        if (ims.shape[-1] != 3):
            ims = ims.transpose((0,2,3,1))
    f = plt.figure(figsize=figsize)
    cols = len(ims)//rows if len(ims) % 2 == 0 else len(ims)//rows + 1
    for i in range(len(ims)):
        sp = f.add_subplot(rows, cols, i+1)
        sp.axis('Off')
        if titles is not None:
            sp.set_title(titles[i], fontsize=16)
        plt.imshow(ims[i], interpolation=None if interp else 'none')

Now, lets print the first batch of training images:

imgs, labels = next(train_batches)
plots(imgs, titles=labels)

Output:

We can see the scaled images of 10 cats and dogs. If you run again the above code, it will fetch next 10 images from training dataset as we are using batch size of 10 for training images.

Step 5: Load and analyze VGG16 model

vgg16_model = keras.applications.vgg16.VGG16()
vgg16_model.summary()
type(vgg16_model)

In the above code, first line will load the VGG16 model. It may take some time. By executing second line, we can see summary of the existing model. It has a lot of convolutional, pooling and dense layers. Executing third line, we can see this model is of type "Model". In next step, we will create a model of type "Sequential".

Step 6: Fine-tune VGG16 model

Following are the steps involved in fine-tuning a model:

1. Copy all the hidden layers in a new model
2. Remove output layer
3. Freeze the hidden layers
4. Add custom output layer

For more details on fine-tuning a model, please visit my this post.

Lets perform all the above steps.

model = Sequential() for layer in vgg16_model.layers[:-1]: model.add(layer)

In the above code, we have created a new sequential model and copied all the layers of VGG16 model except the last layer which is an output layer. We have done this because we want our custom output layer which will have only two nodes as our image classification problem has only two classes (cats and dogs).

Now, if we execute following statement, we will get replica of existing VGG16 model, except output layer.

model.summary()

Now, lets freeze the hidden layers as we don't want to change any weight and bias associated with these layers. We want to use these layers as it is as all these layers are already well trained on image classification problem.

for layer in model.layers: layer.trainable = False

Now, add a custom output layer with only two nodes and softmax as activation function.

model.add(Dense(2, activation='softmax'))
model.summary()

Now, our new fine-tuned model is ready. Lets train it with new data and then predict from it.

Step 7: Compile the model

model.compile(Adam(lr=0.0001), loss='categorical_crossentropy', metrics=['accuracy'])

Using Adam as an optimizer and categorical cross entropy as loss function.

Step 8: Train the model

model.fit_generator(train_batches, steps_per_epoch=4, validation_data=valid_batches, validation_steps=4, epochs=5, verbose=2)

Executing this step will take some time as we are using 5 epochs.

Step 9: Predict from the model

Lets print first batch of the test images.

test_imgs, test_labels = next(test_batches) plots(test_imgs, titles=test_labels)

From the output, we can see that it shows the final results in form of [0. 1.], [1. 0.] etc. Lets format this output so that we can get it in form of 0, 1 etc.

test_labels = test_labels[:,0] test_labels

Now, finally make prediction.

predictions = model.predict_generator(test_batches, steps=1, verbose=0)
predictions

It shows the predictions in form of probabilities. Lets round it off.

rounded_predictions = np.round(predictions[:,0])
rounded_predictions

Step 10: Check the accuracy

confusionMatrix = confusion_matrix(test_labels, rounded_predictions)
accuracyScore = accuracy_score(test_labels, rounded_predictions)
classificationReport = classification_report(test_labels, rounded_predictions)
print(confusionMatrix)
print(accuracyScore * 100)
print(classificationReport)

Please note that we won't get desired accuracy with this small dataset. We need thousands of image to train our model to get desired accuracy. We can use data augmentation to increase the data. You can download thousands of images of cats and dogs from Kaggle to train this model.