keras_resnet50.py

# -*- coding: utf-8 -*-
"""Keras_ResNet50.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1o_mZbTHbiEVeWmWS0n1eF7XAUfZ-jIzf
"""

#Import necessary modules
import tensorflow as tf
from tensorflow.keras.layers import GlobalAveragePooling2D, Dense
from tensorflow.keras.models import Model
from tensorflow.keras.applications.resnet50 import ResNet50
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard

train_dir = 'train'
test_dir = 'test'
num_train = 1000 #number of train images
num_test = 260   #number of test images
batch_size = 64
num_layers = 91  #number of lower layers not to be trained
save_path='models/resnet50'
log_dir='./resnet50_logs'
fine_tuning = True

#Loading pre-trained ResNet50 model 
base = ResNet50(weights='imagenet', include_top=False, input_shape=(224,224,3))

#Intializing top of the model
x = base.output

#Applying global average pooling to the loaded output of vgg16 model
x = GlobalAveragePooling2D()(x)
x = Dense(1024, activation='relu')(x)
model_output = Dense(10, activation='softmax')(x)

#Instantiation of Model functional API with input as base model input and model_output as output
model = Model(inputs=base.input, outputs=model_output)

#Compiling the model
model.compile(optimizer='Adam', loss='categorical_crossentropy', metrics=['accuracy'])

#Instantiation of ImageDataGenerator class
train_generator = ImageDataGenerator(
                                      rescale=1./255,
                                      shear_range=0.2,
                                      zoom_range=0.2,
                                      horizontal_flip=True)

test_generator = ImageDataGenerator(rescale=1./255)

#Takig the path to a directory & generating batches of augmented data
train_flow = train_generator.flow_from_directory(
                                                  train_dir, 
                                                  target_size=(224, 224),
                                                  batch_size=batch_size,
                                                  class_mode='categorical')

test_flow = test_generator.flow_from_directory(
                                                  test_dir, 
                                                  target_size=(224, 224),
                                                  batch_size=batch_size,
                                                  class_mode='categorical')

#Whether to perform fine-tuning or transfer learning
if fine_tuning:

  #Freezing the lower layer and making the rest trainable for fine-tuning                                            
  for layer in model.layers[:num_layers]:
    layer.trainable = False

  for layer in model.layers[num_layers:]:
    layer.trainable = True
else:
  #Train only top layer for transer learning
  for layer in base.layers:
    layer.trainable = False
  
#Saving the model on evey iteration
model_check_point = ModelCheckpoint(save_path, monitor='val_loss', verbose=0)

#Tensorboard to visualize the results
tensorboard = TensorBoard(log_dir=log_dir)

#Training the model on data generated batch-by-batch
model.fit_generator(
                      train_flow,
                      steps_per_epoch=num_train//batch_size,
                      epochs=10,
                      validation_data=test_flow,
                      validation_steps=num_test//batch_size, 
                      callbacks=[tensorboard, model_check_point])