Data augmentation for images in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Data augmentation for images

Problem:You have a small image dataset to train a classifier. The model fits training data well with 98% accuracy but validation accuracy is only 70%. This shows overfitting due to limited data variety.

Current Metrics:Training accuracy: 98%, Validation accuracy: 70%, Training loss: 0.05, Validation loss: 0.85

Issue:The model overfits because it memorizes training images but cannot generalize to new images. The dataset is too small and lacks diversity.

Your Task

Use data augmentation to increase image variety and reduce overfitting. Target validation accuracy >85% while keeping training accuracy below 92%.

Do not change the model architecture.

Do not add more original images.

Only modify the data pipeline to include augmentation.

Hint 1

Hint 2

Hint 3

Solution

TensorFlow

import tensorflow as tf
from tensorflow.keras import layers, models

# Load example dataset
(train_images, train_labels), (val_images, val_labels) = tf.keras.datasets.cifar10.load_data()

# Normalize images
train_images = train_images.astype('float32') / 255.0
val_images = val_images.astype('float32') / 255.0

# Define data augmentation pipeline
data_augmentation = tf.keras.Sequential([
    layers.RandomFlip('horizontal'),
    layers.RandomRotation(0.1),
    layers.RandomZoom(0.1),
    layers.RandomTranslation(0.1, 0.1)
])

# Build model (simple CNN)
model = models.Sequential([
    layers.Input(shape=(32, 32, 3)),
    data_augmentation,  # Apply augmentation only during training
    layers.Conv2D(32, 3, activation='relu'),
    layers.MaxPooling2D(),
    layers.Conv2D(64, 3, activation='relu'),
    layers.MaxPooling2D(),
    layers.Flatten(),
    layers.Dense(64, activation='relu'),
    layers.Dense(10, activation='softmax')
])

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

# Train model
history = model.fit(train_images, train_labels, epochs=20, batch_size=64, validation_data=(val_images, val_labels))

Added a data augmentation pipeline using RandomFlip, RandomRotation, RandomZoom, and RandomTranslation layers.

Inserted the augmentation pipeline as the first layer in the model to augment training images on the fly.

Kept validation data unchanged to fairly evaluate model generalization.

Results Interpretation

Before augmentation: Training accuracy was 98%, validation accuracy was 70%. The large gap showed overfitting.

After augmentation: Training accuracy dropped to 90%, validation accuracy improved to 87%. Loss values also show better generalization.

Data augmentation artificially increases training data diversity. This helps the model learn more general features and reduces overfitting, improving validation accuracy.

Bonus Experiment

Try adding dropout layers to the model along with data augmentation to further reduce overfitting.

💡 Hint

Insert dropout layers with rates like 0.3 after convolutional or dense layers to randomly deactivate neurons during training.