Callbacks (EarlyStopping, ModelCheckpoint) in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Callbacks (EarlyStopping, ModelCheckpoint)

Problem:Train a neural network to classify handwritten digits using the MNIST dataset.

Current Metrics:Training accuracy: 99%, Validation accuracy: 85%, Validation loss fluctuates and sometimes increases after some epochs.

Issue:The model overfits: training accuracy is very high but validation accuracy is much lower and validation loss does not improve steadily.

Your Task

Use EarlyStopping and ModelCheckpoint callbacks to reduce overfitting and improve validation accuracy to at least 90%.

You must keep the same model architecture.

You can only add or modify callbacks and training parameters like epochs and batch size.

Do not change the dataset or preprocessing.

Hint 1

Hint 2

Hint 3

Solution

TensorFlow

import tensorflow as tf
from tensorflow.keras.datasets import mnist
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint

# Load data
(X_train, y_train), (X_test, y_test) = mnist.load_data()

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

# Build model
model = Sequential([
    Flatten(input_shape=(28, 28)),
    Dense(128, activation='relu'),
    Dense(10, activation='softmax')
])

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

# Callbacks
early_stop = EarlyStopping(monitor='val_loss', patience=3, restore_best_weights=True)
checkpoint = ModelCheckpoint('best_model.h5', monitor='val_loss', save_best_only=True, save_weights_only=False)

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

# Evaluate best model
loss, accuracy = model.evaluate(X_test, y_test)
print(f'Test loss: {loss:.4f}, Test accuracy: {accuracy:.4f}')

Added EarlyStopping callback to stop training when validation loss does not improve for 3 epochs.

Added ModelCheckpoint callback to save the best model based on validation loss.

Reduced maximum epochs to 30 to prevent overtraining.

Results Interpretation

Before: Training accuracy: 99%, Validation accuracy: 85%, Validation loss fluctuates.

After: Training accuracy: 95%, Validation accuracy: 91%, Validation loss decreases steadily.

Using EarlyStopping prevents the model from training too long and overfitting. ModelCheckpoint ensures the best model is saved. Together, they improve validation accuracy and model generalization.

Bonus Experiment

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

💡 Hint

Add a Dropout layer with rate 0.3 after the Dense layer and observe changes in validation accuracy.