Multi-class classification in ML Python - ML Experiment: Train & Evaluate

Experiment - Multi-class classification

Problem:Classify images of handwritten digits (0-9) into 10 classes 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 at least 85% while keeping training accuracy below 92%.

You can only change the model architecture and training hyperparameters.

Do not change the dataset or preprocessing steps.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

ML Python

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='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=0)

# 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 with 50% rate after dense layers to reduce overfitting.

Reduced number of neurons from 128 to 64 and 32 to lower model complexity.

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

Kept optimizer as Adam with default learning rate for stable training.

Results Interpretation

Before: Training accuracy 98%, Validation accuracy 75%, Training loss 0.05, Validation loss 0.85

After: Training accuracy 90%, Validation accuracy 87%, Training loss 0.25, Validation loss 0.35

Adding dropout and early stopping helps reduce overfitting by preventing the model from memorizing training data. This improves validation accuracy and generalization.

Bonus Experiment

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

💡 Hint

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