Input shape specification in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Input shape specification

Problem:You are building a neural network to classify images of size 28x28 pixels in grayscale. The current model does not specify the input shape correctly, causing errors or poor training.

Current Metrics:Model training fails or accuracy is very low due to incorrect input shape.

Issue:The input shape is not properly defined in the first layer of the model, leading to shape mismatch errors or inability to train.

Your Task

Correctly specify the input shape in the model so it matches the data shape (28x28 grayscale images). The model should train successfully and achieve at least 80% accuracy on the test set.

Do not change the dataset or model architecture except for input shape.

Use TensorFlow and Keras only.

Hint 1

Hint 2

Hint 3

Solution

TensorFlow

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

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

# Reshape data to add channel dimension (grayscale = 1 channel)
train_images = train_images.reshape((-1, 28, 28, 1)).astype('float32') / 255.0
test_images = test_images.reshape((-1, 28, 28, 1)).astype('float32') / 255.0

# Build model with correct input shape
model = models.Sequential([
    layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    layers.MaxPooling2D((2, 2)),
    layers.Flatten(),
    layers.Dense(64, activation='relu'),
    layers.Dense(10, activation='softmax')
])

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

# Train model
model.fit(train_images, train_labels, epochs=5, batch_size=64, validation_split=0.2)

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

Added input_shape=(28, 28, 1) to the first Conv2D layer to match grayscale image shape.

Reshaped training and test images to have shape (num_samples, 28, 28, 1).

Normalized pixel values to range 0-1 by dividing by 255.

Results Interpretation

Before: Model training failed or accuracy was very low due to input shape mismatch.

After: Model trains successfully and achieves about 85% accuracy on test data.

Specifying the correct input shape is essential for the model to understand the data format and train properly. For image data, the input shape must include height, width, and channels.

Bonus Experiment

Try changing the input shape to (28, 28) without the channel dimension and observe what error or behavior occurs.

💡 Hint

TensorFlow Conv2D layers expect 3D input (height, width, channels). Omitting channels causes shape mismatch errors.