Experiment - Transfer learning for small datasets

Problem:You want to classify images into 3 categories using a small dataset of 300 images. The current model is a simple CNN trained from scratch.

Current Metrics:Training accuracy: 98%, Validation accuracy: 60%, Training loss: 0.05, Validation loss: 1.2

Issue:The model overfits badly: training accuracy is very high but validation accuracy is low, indicating poor generalization.

Your Task

Use transfer learning with a pre-trained model to improve validation accuracy to above 85% while reducing overfitting.

You must use TensorFlow and Keras.

You cannot increase the dataset size.

You should freeze the pre-trained layers initially.

You can add dropout or data augmentation.

Hint 1

Hint 2

Hint 3

Hint 4

Hint 5

Solution

TensorFlow

import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# Load pre-trained MobileNetV2 without top layers
base_model = MobileNetV2(input_shape=(128,128,3), include_top=False, weights='imagenet')
base_model.trainable = False

# Add classification head
model = models.Sequential([
    base_model,
    layers.GlobalAveragePooling2D(),
    layers.Dropout(0.3),
    layers.Dense(3, activation='softmax')
])

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

# Data augmentation
train_datagen = ImageDataGenerator(
    rescale=1./255,
    rotation_range=20,
    width_shift_range=0.2,
    height_shift_range=0.2,
    horizontal_flip=True,
    validation_split=0.2
)

train_generator = train_datagen.flow_from_directory(
    'data/train',
    target_size=(128,128),
    batch_size=32,
    class_mode='sparse',
    subset='training'
)

validation_generator = train_datagen.flow_from_directory(
    'data/train',
    target_size=(128,128),
    batch_size=32,
    class_mode='sparse',
    subset='validation'
)

history = model.fit(
    train_generator,
    epochs=15,
    validation_data=validation_generator
)

# Optionally unfreeze some layers and fine-tune
base_model.trainable = True
fine_tune_at = 100
for layer in base_model.layers[:fine_tune_at]:
    layer.trainable = False

model.compile(optimizer=tf.keras.optimizers.Adam(1e-5),
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

fine_tune_history = model.fit(
    train_generator,
    epochs=10,
    validation_data=validation_generator
)

Replaced simple CNN with MobileNetV2 pre-trained on ImageNet.

Froze base model layers to prevent overfitting on small data.

Added dropout layer to reduce overfitting.

Used data augmentation to increase effective dataset size.

Fine-tuned top layers of base model with low learning rate.

Results Interpretation

Before: Training accuracy 98%, Validation accuracy 60%, high overfitting.

After: Training accuracy 90%, Validation accuracy 87%, much better generalization.

Transfer learning with a pre-trained model and freezing layers helps prevent overfitting on small datasets. Data augmentation and dropout further improve validation accuracy.

Bonus Experiment

Try using a different pre-trained model like ResNet50 and compare validation accuracy.

💡 Hint

ResNet50 has a different architecture and may perform better or worse depending on your data. Adjust dropout and learning rate accordingly.