Batch size and epochs in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Batch size and epochs

Problem:Train a neural network to classify handwritten digits using 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 by adjusting batch size and number of epochs to improve validation accuracy to above 90% while keeping training accuracy below 95%.

You can only change batch size and number of epochs.

Do not change the model architecture or optimizer.

Hint 1

Hint 2

Hint 3

Solution

TensorFlow

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

# Load MNIST data
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()

# Normalize images
train_images = train_images / 255.0
test_images = test_images / 255.0

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

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

# Train model with adjusted batch size and epochs
history = model.fit(train_images, train_labels,
                    epochs=10,
                    batch_size=128,
                    validation_split=0.2,
                    verbose=2)

# Evaluate on test data
test_loss, test_acc = model.evaluate(test_images, test_labels, verbose=0)
print(f'Test accuracy: {test_acc:.4f}')

Increased batch size from 32 to 128 to reduce noise in gradient updates.

Reduced epochs from 20 to 10 to prevent overfitting.

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.18, Validation loss: 0.25

Adjusting batch size and epochs can reduce overfitting by controlling how the model learns. Larger batch sizes provide smoother updates, and fewer epochs prevent the model from memorizing training data.

Bonus Experiment

Try adding dropout layers to the model to further reduce overfitting and improve validation accuracy.

💡 Hint

Add a dropout layer with rate 0.3 after the first Dense layer and retrain with the same batch size and epochs.