model.fit() training loop in TensorFlow - ML Experiment: Train & Evaluate

Experiment - model.fit() training loop

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

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

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 90% while keeping training accuracy below 95%.

You can only modify the model.fit() training loop parameters and add callbacks.

Do not change the model architecture or dataset.

Hint 1

Hint 2

Hint 3

Solution

TensorFlow

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

# 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 simple model
model = models.Sequential([
    layers.Dense(128, activation='relu', input_shape=(28*28,)),
    layers.Dense(10, activation='softmax')
])

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

# Add early stopping callback
early_stop = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=3, restore_best_weights=True)

# Train model with validation split and early stopping
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:.4f}, Test loss: {test_loss:.4f}')

Added validation_split=0.2 to monitor validation performance during training.

Added EarlyStopping callback to stop training early when validation loss stops improving.

Reduced batch size to 64 for better generalization.

Increased max epochs to 30 but training stops early due to early stopping.

Results Interpretation

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

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

Using validation data during training and stopping early when validation loss stops improving helps reduce overfitting and improves the model's ability to generalize to new data.

Bonus Experiment

Try adding dropout layers to the model to further reduce overfitting and compare the results.

💡 Hint

Add a Dropout layer with rate 0.3 after the first Dense layer and retrain the model with the same training loop.