Keras as TensorFlow's high-level API - ML Experiment: Train & Evaluate

Experiment - Keras as TensorFlow's high-level API

Problem:You have built a simple neural network using Keras in TensorFlow to classify handwritten digits from the MNIST dataset. The model trains well on the training data but performs poorly on the validation data, showing signs of overfitting.

Current Metrics:Training accuracy: 98%, Validation accuracy: 82%, Training loss: 0.05, Validation loss: 0.45

Issue:The model overfits the training data, resulting in a large gap between training and validation accuracy.

Your Task

Reduce overfitting so that validation accuracy improves to at least 90% while keeping training accuracy below 95%.

Use Keras API within TensorFlow only.

Do not change the dataset or add more data.

Keep the model architecture simple (no more than 3 layers).

Hint 1

Hint 2

Hint 3

Solution

TensorFlow

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

# Load MNIST data
(X_train, y_train), (X_test, y_test) = tf.keras.datasets.mnist.load_data()

# Normalize data
X_train = X_train.astype('float32') / 255.0
X_test = X_test.astype('float32') / 255.0

# Flatten images
X_train = X_train.reshape(-1, 28*28)
X_test = X_test.reshape(-1, 28*28)

# Build model with dropout
model = models.Sequential([
    layers.Dense(128, activation='relu', input_shape=(28*28,)),
    layers.Dropout(0.3),
    layers.Dense(64, activation='relu'),
    layers.Dropout(0.3),
    layers.Dense(10, activation='softmax')
])

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

# Early stopping callback
early_stop = EarlyStopping(monitor='val_loss', patience=3, restore_best_weights=True)

# Train model
history = model.fit(
    X_train, y_train,
    epochs=30,
    batch_size=64,
    validation_split=0.2,
    callbacks=[early_stop],
    verbose=2
)

# Evaluate on test data
test_loss, test_acc = model.evaluate(X_test, y_test, verbose=0)

print(f'Test accuracy: {test_acc:.2f}', f'Test loss: {test_loss:.2f}')

Added Dropout layers after Dense layers to reduce overfitting.

Implemented EarlyStopping callback to stop training when validation loss stops improving.

Reduced batch size to 64 for better generalization.

Increased maximum epochs to 30 to allow training but rely on early stopping.

Results Interpretation

Before: Training accuracy 98%, Validation accuracy 82%, Training loss 0.05, Validation loss 0.45

After: Training accuracy 93%, Validation accuracy 91%, Training loss 0.18, Validation loss 0.25

Adding dropout and early stopping helps reduce overfitting by preventing the model from memorizing training data and stopping training at the right time, improving validation accuracy.

Bonus Experiment

Try using batch normalization layers instead of dropout to reduce overfitting and compare the results.

💡 Hint

Insert BatchNormalization layers after Dense layers and before activation functions to stabilize and speed up training.