Transfer Learning for CIFAR-10 Classification Using VGG16

Abstract

In this experiment, I trained a convolutional neural network (CNN) using transfer learning to classify images from the CIFAR-10 dataset. By leveraging the VGG16 architecture pre-trained on ImageNet, I aimed to achieve a validation accuracy of 87% or higher. The model was trained using TensorFlow and Keras, and the final trained model was saved as cifar10.h5.

Introduction

The CIFAR-10 dataset is a popular benchmark dataset in the field of computer vision, consisting of 60,000 32x32 color images in 10 classes. Training a high-accuracy model from scratch on this dataset can be challenging due to its complexity and relatively small image size. Transfer learning, a technique where a model pre-trained on a large dataset is adapted for a different but related task, can significantly improve performance and reduce training time. In this experiment, I utilized the VGG16 model, a deep CNN pre-trained on the ImageNet dataset, to classify CIFAR-10 images.

Materials and Methods

Dataset

The CIFAR-10 dataset was used in this experiment. It contains 60,000 images split into 50,000 training images and 10,000 test images, with 10 classes: airplanes, cars, birds, cats, deer, dogs, frogs, horses, ships, and trucks.

Preprocessing

I implemented a preprocessing function to normalize the pixel values of the images to the range [0, 1] and convert the labels to one-hot encoding.

def preprocess_data(X, Y):
    X_p = X.astype('float32') / 255.0
    Y_p = K.utils.to_categorical(Y, 10)
    return X_p, Y_p
        

Model Architecture

I used the VGG16 model, excluding its top classification layer. I added a lambda layer to resize the CIFAR-10 images from 32x32 to 224x224, followed by the VGG16 model with frozen layers. Finally, I added new trainable layers on top to adapt the model for CIFAR-10 classification.

from tensorflow import keras as K
from tensorflow.keras import layers
from tensorflow.keras.applications import VGG16

def create_model():
    base_model = VGG16(weights='imagenet', include_top=False, input_shape=(224, 224, 3))

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

    model = K.Sequential()
    model.add(layers.Lambda(lambda image: K.backend.resize_images(image, 7, 7, "channels_last"), input_shape=(32, 32, 3)))
    model.add(base_model)
    model.add(layers.Flatten())
    model.add(layers.Dense(256, activation='relu'))
    model.add(layers.Dropout(0.5))
    model.add(layers.Dense(10, activation='softmax'))

    model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
    return model
        

Training and Evaluation

I trained the model on the preprocessed CIFAR-10 training data for 10 epochs with a batch size of 128. The model was then evaluated on the test data to measure its performance.

def main():
    (X_train, Y_train), (X_test, Y_test) = K.datasets.cifar10.load_data()
    X_train, Y_train = preprocess_data(X_train, Y_train)
    X_test, Y_test = preprocess_data(X_test, Y_test)

    model = create_model()
    model.fit(X_train, Y_train, epochs=10, batch_size=128, validation_data=(X_test, Y_test))
    model.save('cifar10.h5')

if __name__ == '__main__':
    main()
        

Results

The model achieved a validation accuracy of over 87% on the CIFAR-10 test set. The use of transfer learning with the VGG16 model significantly improved the accuracy compared to training a CNN from scratch.

Discussion

The experiment demonstrates the effectiveness of transfer learning in improving the performance of image classification tasks on small datasets like CIFAR-10. By leveraging the rich feature representations learned by the VGG16 model on the ImageNet dataset, we were able to achieve high accuracy with relatively few epochs of training.

Acknowledgments

I would like to thank Khaled Haguiga and Santiago Echeverri for their invaluable assistance and support throughout this experiment.

Literature Cited

• Chuanqi Tan, Fuchun Sun, Tao Kong, Wenchang Zhang, Chao Yang, and Chunfang Liu (2018). "A Survey on Deep Transfer Learning" arXiv:1808.01974v1 [cs.LG] 6 Aug 2018. Link
• Towards Data Science. "A Comprehensive Hands-on Guide to Transfer Learning with Real-World Applications in Deep Learning." Link.
• Mukhlif, A. A., Al-Khateeb, B., & Mohammed, M. A. (2023). Incorporating a Novel Dual Transfer Learning Approach for Medical Images. Sensors, 23(2), 570. Link

This blog post summarizes my experimental process of using transfer learning with the VGG16 model to classify the CIFAR-10 dataset. The steps, code snippets, and results presented here provide a clear guide for replicating and understanding the process.