Flatten and Dense layers in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Flatten and Dense layers

Problem:You are building a simple neural network to classify images from the MNIST dataset. The current model uses Flatten and Dense layers but is not learning well.

Current Metrics:Training accuracy: 60%, Validation accuracy: 58%, Training loss: 1.2, Validation loss: 1.3

Issue:The model is underfitting. It is not learning enough from the data, resulting in low accuracy on both training and validation sets.

Your Task

Improve the model so that training accuracy reaches at least 90% and validation accuracy reaches at least 88%.

You must keep using Flatten and Dense layers only.

You cannot add convolutional layers or change the dataset.

You can adjust the number of neurons, activation functions, and training parameters.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

TensorFlow

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

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

# Normalize pixel values to [0,1]
X_train = X_train / 255.0
X_test = X_test / 255.0

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

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

# Train the model
history = model.fit(X_train, y_train, epochs=10, batch_size=32, validation_split=0.2)

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

print(f'Test accuracy: {test_acc:.4f}')

Normalized input images by dividing pixel values by 255.

Increased number of neurons in Dense layers to 128 and 64.

Added ReLU activation to hidden Dense layers and softmax to output layer.

Used Adam optimizer for better training.

Increased training epochs to 10 and batch size to 32.

Results Interpretation

Before: Training accuracy 60%, Validation accuracy 58%, Loss around 1.2-1.3

After: Training accuracy 95%, Validation accuracy 92%, Loss reduced to about 0.15-0.20

Normalizing input data and increasing model capacity with more neurons and proper activations help the model learn better. Using suitable optimizer and training longer improves accuracy significantly.

Bonus Experiment

Try adding dropout layers between Dense layers to reduce overfitting and see if validation accuracy improves further.

💡 Hint

Add layers.Dropout(0.3) after Dense layers and retrain the model.