Emotion-Classification/code/model.py

import json
import pickle
import numpy as np
import tensorflow as tf
from tensorflow.keras import optimizers
from tensorflow.keras.optimizers.experimental import SGD
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Flatten, Dense, Conv2D, GlobalAveragePooling2D
from tensorflow.keras.layers import Dropout, BatchNormalization, Activation
from tensorflow.keras.callbacks import Callback, EarlyStopping, ReduceLROnPlateau
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras import regularizers
from sklearn.utils.class_weight import compute_class_weight


def build_model(num_classes):
    # Load pre-trained model
    base_model = tf.keras.applications.VGG19(weights='imagenet', include_top=False, input_shape=(48, 48, 3))
    
    # Add dense layers
    x = base_model.layers[-2].output
    x = GlobalAveragePooling2D()(x)
    
    # Add final classification layer
    output_layer = Dense(num_classes, activation='softmax')(x)
    
    # Create model
    model = Model(inputs=base_model.input, outputs=output_layer)
    
    # Compile model with hyperparameters
    opt = SGD(learning_rate=0.0092,momentum=0.9,nesterov=True)
        
    model.compile(loss='categorical_crossentropy', optimizer=opt, metrics=['accuracy'])
    
    return model


def train(model, X_train, y_train, X_valid, y_valid):
    params = {}
    #preserve pre-trained model weights and make layers non-trainable 

    lr_scheduler = ReduceLROnPlateau(monitor = 'val_accuracy', 
                              factor = 0.4, 
                              patience = 7,
                              min_lr = 1e-6,
                              verbose = 1)

    early_stopping = EarlyStopping(monitor = 'val_accuracy', 
                                min_delta = 0.00005, 
                                patience = 11,
                                verbose = 1, 
                                restore_best_weights = True)

    callbacks = [early_stopping,lr_scheduler]

# Define image data generator
    train_datagen = ImageDataGenerator(rotation_range=20,
                                    width_shift_range=0.20,
                                    height_shift_range=0.20,
                                    shear_range=0.15,
                                    zoom_range=0.15,
                                    horizontal_flip=True,
                                    fill_mode='nearest')
    train_datagen.fit(X_train)

    batch_size = 96
    epochs = 24
    steps_per_epoch = (int)(len(X_train) / batch_size)
    optims = [SGD(learning_rate = 0.0092, momentum=0.9, nesterov=True)]

    model.compile(loss = 'categorical_crossentropy',
            optimizer = optims[0],
            metrics = ['accuracy'])
    
    class_weights = compute_class_weight('balanced', np.unique(y_train), y_train)
    class_weights_dict = dict(enumerate(class_weights)

    history = model.fit(train_datagen.flow(X_train, 
                                       y_train, 
                                       batch_size = batch_size),
                                       validation_data = (X_valid, y_valid),
                                       steps_per_epoch = steps_per_epoch,
                                       epochs = epochs,
                                       callbacks = callbacks,
                                       use_multiprocessing = False,
                                       class_weight=class_weights_dict)
    
    model.save('data/model.h5')

    return history


def log_metrics(history):
    from dagshub import dagshub_logger, DAGsHubLogger

    # Option 1 - As a context manager:
    with dagshub_logger(metrics_path="logs/metrics.csv", hparams_path="logs/params.yml") as logger:

        # Metric logging: 
        logger.log_metrics(train_accuracy=history.history['accuracy'][-1], step_num=1) 
        logger.log_metrics(val_accuracy=history.history['val_accuracy'][-1], step_num=1) 
        logger.log_metrics(train_loss=history.history['loss'][-1], step_num=1) 
        logger.log_metrics(val_loss=history.history['val_loss'][-1], step_num=1) 

        # Hyperparameters logging: 
        logger.log_hyperparams(optimizer={'type': 'sgd', 'learning_rate': 0.0092, 'momentum': 0.90, 'nesterov': True})
        logger.log_hyperparams(loss='categorical_crossentropy')
        logger.log_hyperparams(batch_size=96)
        logger.log_hyperparams(steps_per_epoch=(int)(len(X_train) / 96))
        logger.log_hyperparams(epochs = 24)
        logger.log_hyperparams(callbacks=['EarlyStopping', 'ReduceLROnPlateau'])
        logger.log_hyperparams(data_augmentation= 'ImageDataGenerator')




if __name__ == '__main__':

    X_train = np.load('data/processed/X_train.npy')
    # load data from y_train.npy in data/preprocessed directory
    y_train = np.load('data/processed/y_train.npy')
    # load data from X_test.npy in data/preprocessed directory
    X_valid = np.load('data/processed/X_valid.npy')
    # load data from y_test.npy in data/preprocessed directory
    y_valid = np.load('data/processed/y_valid.npy')

    img_width = X_train.shape[1]
    img_height = X_train.shape[2]
    img_depth = X_train.shape[3]
    num_classes = y_train.shape[1] 

    import os
    from dotenv import load_dotenv

    # Load environment variables from .env file
    load_dotenv()

    # Access environment variables
    USERNAME = os.getenv('MLFLOW_TRACKING_USERNAME')
    PASSWORD = os.getenv('MLFLOW_TRACKING_PASSWORD')

    import mlflow
    mlflow.set_tracking_uri("https://dagshub.com/GauravMohan1/Emotion-Classification.mlflow")
    with mlflow.start_run():
        # batch size of 32 performs the best.
        model = build_model(num_classes)

        history = train(model, X_train, y_train, X_valid, y_valid)
        metrics = {"train_accuracy": history.history['accuracy'][-1], "val_accuracy": history.history['val_accuracy'][-1], 
                   "train_loss": history.history['loss'][-1], "val_loss": history.history['val_loss'][-1]}
        
        params = {"optimizer": {'type': 'sgd', 'learning_rate': 0.0092, 'momentum': 0.90, 'nesterov': True}, "loss": 'categorical_crossentropy', 'batch_size': 96,
                  'epochs': 24, 'callbacks': ['EarlyStopping', 'ReduceLROnPlateau'], 'data_augmentation': 'ImageDataGenerator'}

        mlflow.log_metrics(metrics)
        mlflow.log_params(params)