Multi-class classification model in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Multi-class classification model

Problem:Classify images of handwritten digits (0-9) using a neural network.

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

Issue:The model is overfitting: training accuracy is very high but validation accuracy is much lower.

Your Task

Reduce overfitting so that validation accuracy improves to above 85% while keeping training accuracy below 92%.

You can only modify the model architecture and training hyperparameters.

Do not change the dataset or preprocessing steps.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

TensorFlow

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

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

# Normalize data
X_train, X_test = X_train / 255.0, X_test / 255.0

# Build model with dropout to reduce overfitting
model = models.Sequential([
    layers.Flatten(input_shape=(28, 28)),
    layers.Dense(64, activation='relu'),
    layers.Dropout(0.5),
    layers.Dense(32, activation='relu'),
    layers.Dropout(0.5),
    layers.Dense(10, activation='softmax')
])

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

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

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

# Evaluate on test data
test_loss, test_accuracy = model.evaluate(X_test, y_test)

Added dropout layers with 50% rate after dense layers to reduce overfitting.

Reduced the number of neurons from 128 and 64 to 64 and 32 to simplify the model.

Added early stopping to stop training when validation loss stops improving.

Kept learning rate at 0.001 for stable training.

Results Interpretation

Before: Training accuracy was 98% but validation accuracy was only 75%, showing overfitting.

After: Training accuracy dropped to 90% but validation accuracy improved to 87%, showing better generalization.

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

Bonus Experiment

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

💡 Hint

Insert batch normalization layers after dense layers and remove dropout layers. Train with the same settings and observe validation accuracy.