vgg_from_scratch_image_classification.py

# import necessary layers  
from tensorflow.keras.layers import Input, Conv2D , Dropout, MaxPool2D, Flatten, Dense 
from tensorflow.keras import Model
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.regularizers import l2
import tensorflow as tf
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
import os
import matplotlib.pyplot as plt
import sys
from tensorflow.keras.callbacks import CSVLogger


MODEL_FNAME = "trained_model.h5"
base_dir = "dataset"
tmp_model_name = "tmp.h5"
INPUT_SIZE = 224
BATCH_SIZE = 16


physical_devices = tf.config.list_physical_devices()
print("DEVICES : \n", physical_devices)


print('Using:')
print('\t\u2022 Python version:',sys.version)
print('\t\u2022 TensorFlow version:', tf.__version__)
print('\t\u2022 tf.keras version:', tf.keras.__version__)
print('\t\u2022 Running on GPU' if tf.test.is_gpu_available() else '\t\u2022 GPU device not found. Running on CPU')

count = 0
previous_acc = 0


if not os.path.exists(MODEL_FNAME):
    
    """ Create VGG Model"""
    
    # input
    input = Input(shape =(INPUT_SIZE,INPUT_SIZE,3))
    
    weight_initializer = tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.01, seed=None)
    bias_initializer=tf.keras.initializers.Zeros()
    
    # 1st Conv Block
    
    x = Conv2D (filters =64, kernel_size =3, padding ='same', activation='relu',kernel_initializer=weight_initializer,kernel_regularizer=l2(0.00005),bias_initializer=bias_initializer)(input)
    x = Conv2D (filters =64, kernel_size =3, padding ='same', activation='relu',kernel_initializer=weight_initializer,kernel_regularizer=l2(0.00005),bias_initializer=bias_initializer)(x)
    x = MaxPool2D(pool_size =2, strides =2, padding ='same')(x)
    
    # 2nd Conv Block
    
    x = Conv2D (filters =128, kernel_size =3, padding ='same', activation='relu',kernel_initializer=weight_initializer,kernel_regularizer=l2(0.00005),bias_initializer=bias_initializer)(x)
    x = Conv2D (filters =128, kernel_size =3, padding ='same', activation='relu',kernel_initializer=weight_initializer,kernel_regularizer=l2(0.00005),bias_initializer=bias_initializer)(x)
    x = MaxPool2D(pool_size =2, strides =2, padding ='same')(x)
    
    # 3rd Conv block
    
    x = Conv2D (filters =256, kernel_size =3, padding ='same', activation='relu',kernel_initializer=weight_initializer,kernel_regularizer=l2(0.00005),bias_initializer=bias_initializer)(x)
    x = Conv2D (filters =256, kernel_size =3, padding ='same', activation='relu',kernel_initializer=weight_initializer,kernel_regularizer=l2(0.00005),bias_initializer=bias_initializer)(x)
    x = Conv2D (filters =256, kernel_size =3, padding ='same', activation='relu',kernel_initializer=weight_initializer,kernel_regularizer=l2(0.00005),bias_initializer=bias_initializer)(x)
    x = MaxPool2D(pool_size =2, strides =2, padding ='same')(x)
    
    # 4th Conv block
    
    x = Conv2D (filters =512, kernel_size =3, padding ='same', activation='relu',kernel_initializer=weight_initializer,kernel_regularizer=l2(0.00005),bias_initializer=bias_initializer)(x)
    x = Conv2D (filters =512, kernel_size =3, padding ='same', activation='relu',kernel_initializer=weight_initializer,kernel_regularizer=l2(0.00005),bias_initializer=bias_initializer)(x)
    x = Conv2D (filters =512, kernel_size =3, padding ='same', activation='relu',kernel_initializer=weight_initializer,kernel_regularizer=l2(0.00005),bias_initializer=bias_initializer)(x)
    x = MaxPool2D(pool_size =2, strides =2, padding ='same')(x)
    
    
    # 5th Conv block
    
    x = Conv2D (filters =512, kernel_size =3, padding ='same', activation='relu',kernel_initializer=weight_initializer,kernel_regularizer=l2(0.00005),bias_initializer=bias_initializer)(x)
    x = Conv2D (filters =512, kernel_size =3, padding ='same', activation='relu',kernel_initializer=weight_initializer,kernel_regularizer=l2(0.00005),bias_initializer=bias_initializer)(x)
    x = Conv2D (filters =512, kernel_size =3, padding ='same', activation='relu',kernel_initializer=weight_initializer,kernel_regularizer=l2(0.00005),bias_initializer=bias_initializer)(x)
    x = MaxPool2D(pool_size =2, strides =2, padding ='same')(x)
    
    # Fully connected layers
    
    x = Flatten()(x)
    x = Dropout(0.5)(x)
    x = Dense(units = 4096, activation ='relu', kernel_initializer=weight_initializer,kernel_regularizer=l2(0.00005),bias_initializer=bias_initializer)(x)
    x = Dropout(0.5)(x)
    x = Dense(units = 4096, activation ='relu', kernel_initializer=weight_initializer,kernel_regularizer=l2(0.00005),bias_initializer=bias_initializer)(x)
    output = Dense(units = 2, activation ='softmax')(x)
    
    # creating the model
    
    model = Model (inputs=input, outputs =output)
    
    m = model
    m.save(tmp_model_name)
    del m
    tf.keras.backend.clear_session()
    
    
    model.summary()
    
    """ Prepare the Dataset for Training"""
    
    
    train_dir = os.path.join(base_dir, 'train')
    val_dir = os.path.join(base_dir, 'validation')
 
    
    
    train_batches = ImageDataGenerator(rescale = 1 / 255.).flow_from_directory(train_dir,
                                                         target_size=(INPUT_SIZE,INPUT_SIZE),
                                                         shuffle=True,
                                                         seed=42,
                                                         batch_size=BATCH_SIZE)

    val_batches = ImageDataGenerator(rescale = 1 / 255.).flow_from_directory(val_dir,
                                                         target_size=(INPUT_SIZE,INPUT_SIZE),
                                                         shuffle=True,
                                                         seed=42,
                                                         batch_size=BATCH_SIZE)
    
    """ Train """
    
    class CustomLearningRateScheduler(tf.keras.callbacks.Callback):
        def __init__(self, schedule):
            super(CustomLearningRateScheduler, self).__init__()
            self.schedule = schedule
        

        def on_epoch_end(self, epoch, logs=None):
            if not hasattr(self.model.optimizer, "lr"):
                raise ValueError('Optimizer must have a "lr" attribute.')
            # Get the current learning rate from model's optimizer.
            lr = float(tf.keras.backend.get_value(self.model.optimizer.learning_rate))
            # Call schedule function to get the scheduled learning rate.
            # keys = list(logs.keys())
            # print("keys",keys)
            val_acc = logs.get("val_binary_accuracy")
            scheduled_lr = self.schedule(lr, val_acc)
            # Set the value back to the optimizer before this epoch starts
            tf.keras.backend.set_value(self.model.optimizer.lr, scheduled_lr)
          
    def learning_rate_scheduler(lr, val_acc):
        global count
        global previous_acc
      
        if val_acc == previous_acc:
          #  print("acc ", val_acc, "previous acc ", previous_acc)
            count += 1
        else:
            count = 0
        
        if count >= 5:
            print("acc is the same for 10 epoch, learnin rate decreased by /10")
            count = 0
            lr /= 10
            print("new learning rate:", lr)
            
        previous_acc = val_acc   
        return lr
    
    #compile the model by determining loss function Binary Cross Entropy, optimizer as SGD
    model.compile(optimizer=tf.keras.optimizers.SGD(lr=0.0000001, momentum=0.9),
                  loss=tf.keras.losses.BinaryCrossentropy(),
                  metrics=[tf.keras.metrics.BinaryAccuracy()],
                  sample_weight_mode=[None])
    
    early_stopping = EarlyStopping(monitor='val_loss', patience=10)
        
    checkpointer = ModelCheckpoint(filepath=MODEL_FNAME, verbose=1, save_best_only=True)
    
    csv_logger = CSVLogger('log.csv', append=True, separator=' ')
    
    history=model.fit(train_batches,
        validation_data = val_batches,
        epochs = 100,
        verbose = 1,
        shuffle = True,
        callbacks = [checkpointer,early_stopping,CustomLearningRateScheduler(learning_rate_scheduler),csv_logger])
    
    """ Plot the train and validation 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', 'validation'], loc='upper left')
    plt.show()
    
    """ Plot the train and validation Accuracy """
  
    plt.plot(history.history['binary_accuracy'])
    plt.plot(history.history['val_binary_accuracy'])
    plt.title('model accuracy')
    plt.ylabel('accuracy')
    plt.xlabel('epoch')
    plt.legend(['train', 'validation'], loc='upper left')
    plt.show()
  
    print("End of Training")
    
else: 
    
    """ Test """
    
    test_dir = os.path.join(base_dir, 'test')
    
    test_batches = ImageDataGenerator(rescale = 1 / 255.).flow_from_directory(test_dir,
                                                       target_size=(INPUT_SIZE,INPUT_SIZE),
                                                       class_mode='categorical',
                                                       shuffle=False,
                                                       seed=42,
                                                       batch_size=1)
    model = tf.keras.models.load_model(MODEL_FNAME)
    model.summary()
    
    # Evaluate on test data
    scores = model.evaluate(test_batches)
    print("metric names",model.metrics_names) 
  
    print(model.metrics_names[0], scores[0])
    print(model.metrics_names[1], scores[1])
        

tf.keras.backend.clear_session()