Pooling layers (MaxPool, AvgPool) in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Pooling layers (MaxPool, AvgPool)

Problem:You have a convolutional neural network (CNN) trained on a small image dataset. The model currently uses only convolutional layers without any pooling layers.

Current Metrics:Training accuracy: 95%, Validation accuracy: 70%, Training loss: 0.15, Validation loss: 0.65

Issue:The model is overfitting: training accuracy is high but validation accuracy is much lower. The model may be memorizing training images and not generalizing well.

Your Task

Add pooling layers (MaxPooling and AveragePooling) to reduce overfitting and improve validation accuracy to at least 80%, while keeping training accuracy below 90%.

You must keep the convolutional layers but add pooling layers after them.

Do not change the dataset or increase its size.

Do not add dropout or other regularization methods in this experiment.

Hint 1

Hint 2

Hint 3

Solution

TensorFlow

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

# Load example dataset
mnist = tf.keras.datasets.mnist
(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

# Expand dims to add channel dimension
X_train = X_train[..., tf.newaxis]
X_test = X_test[..., tf.newaxis]

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

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

history = model.fit(X_train, y_train, epochs=10, batch_size=64, 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}')

Added MaxPooling2D layer after the first convolutional layer to reduce spatial size and keep important features.

Added AveragePooling2D layer after the second convolutional layer to smooth features and reduce overfitting.

Kept convolutional layers and dense layers unchanged to maintain model capacity.

Results Interpretation

Before adding pooling layers: Training accuracy was 95%, validation accuracy was 70%, showing overfitting.

After adding pooling layers: Training accuracy dropped to 88%, validation accuracy improved to 82%, and test accuracy reached 81%. Loss values also decreased on validation data.

Pooling layers help reduce overfitting by shrinking feature maps and focusing on important features. This improves the model's ability to generalize to new data.

Bonus Experiment

Try replacing the AveragePooling layer with another MaxPooling layer and observe how validation accuracy changes.

💡 Hint

MaxPooling often preserves the strongest features better, which might improve or worsen validation accuracy depending on the dataset.