Multi-class classification model in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Multi-class classification model

Problem:Classify images of handwritten digits (0-9) using a neural network.

Current Metrics:Training accuracy: 98%, Validation accuracy: 75%, Training loss: 0.05, Validation loss: 0.85

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 above 85% while keeping training accuracy below 92%.

You can only modify the model architecture and training hyperparameters.

Do not change the dataset or preprocessing steps.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

TensorFlow

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

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

# Normalize data
X_train, X_test = X_train / 255.0, X_test / 255.0

# Build model with dropout to reduce overfitting
model = models.Sequential([
    layers.Flatten(input_shape=(28, 28)),
    layers.Dense(64, activation='relu'),
    layers.Dropout(0.5),
    layers.Dense(32, activation='relu'),
    layers.Dropout(0.5),
    layers.Dense(10, activation='softmax')
])

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# Early stopping callback
early_stop = EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True)

# Train model
history = model.fit(X_train, y_train, epochs=50, batch_size=64, validation_split=0.2, callbacks=[early_stop])

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

Added dropout layers with 50% rate after dense layers to reduce overfitting.

Reduced the number of neurons from 128 and 64 to 64 and 32 to simplify the model.

Added early stopping to stop training when validation loss stops improving.

Kept learning rate at 0.001 for stable training.

Results Interpretation

Before: Training accuracy was 98% but validation accuracy was only 75%, showing overfitting.

After: Training accuracy dropped to 90% but validation accuracy improved to 87%, showing better generalization.

Adding dropout and early stopping helps reduce overfitting by preventing the model from memorizing training data and stopping training at the right time.

Bonus Experiment

Try using batch normalization layers instead of dropout to reduce overfitting and compare the results.

💡 Hint

Insert batch normalization layers after dense layers and remove dropout layers. Train with the same settings and observe validation accuracy.

Practice

1.

What activation function is commonly used in the last layer of a multi-class classification model in TensorFlow?

easy

A. Sigmoid

B. ReLU

C. Softmax

D. Tanh

Solution

1. Step 1: Understand the purpose of the last layer in multi-class classification

   The last layer outputs probabilities for each class, so the activation must convert raw scores to probabilities.

2. Step 2: Identify the activation function that outputs probabilities summing to 1

   Softmax converts logits into probabilities that sum to 1, suitable for multi-class classification.

3. Final Answer:

   Softmax -> Option C

4. Quick Check:

   Softmax = last layer activation [OK]

Hint: Use softmax for multi-class output probabilities [OK]

Common Mistakes:

• Using sigmoid which is for binary classification
• Using ReLU which does not output probabilities
• Using tanh which outputs values between -1 and 1

2.

Which loss function should you use in TensorFlow for a multi-class classification model with integer labels?

easy

A. binary_crossentropy

B. sparse_categorical_crossentropy

C. mean_squared_error

D. hinge

Solution

1. Step 1: Identify the label format

   Labels are integer class IDs, not one-hot encoded vectors.

2. Step 2: Choose loss function matching integer labels for multi-class

   Sparse categorical crossentropy works with integer labels directly, unlike categorical crossentropy which needs one-hot labels.

3. Final Answer:

   sparse_categorical_crossentropy -> Option B

4. Quick Check:

   Integer labels = sparse_categorical_crossentropy [OK]

Hint: Use sparse_categorical_crossentropy for integer class labels [OK]

Common Mistakes:

• Using binary_crossentropy which is for two classes
• Using mean_squared_error which is for regression
• Using hinge loss which is for SVMs

3.

What will be the shape of the output tensor from the last layer of this TensorFlow model for multi-class classification with 4 classes?

model = tf.keras.Sequential([
  tf.keras.layers.Dense(10, activation='relu'),
  tf.keras.layers.Dense(4, activation='softmax')
])
inputs = tf.random.uniform((5, 8))
outputs = model(inputs)
print(outputs.shape)

medium

A. (4, 5)

B. (8, 4)

C. (5, 10)

D. (5, 4)

Solution

1. Step 1: Understand input and output shapes

   Input batch size is 5, each input has 8 features. The last Dense layer outputs 4 units (classes).

2. Step 2: Determine output shape from last layer

   Output shape is (batch_size, number_of_classes) = (5, 4).

3. Final Answer:

   (5, 4) -> Option D

4. Quick Check:

   Batch size 5, classes 4 = (5, 4) [OK]

Hint: Output shape = (batch_size, number_of_classes) [OK]

Common Mistakes:

• Confusing batch size and feature dimensions
• Swapping rows and columns in output shape
• Assuming output shape matches input feature size

4.

Identify the error in this TensorFlow multi-class classification model code:

model = tf.keras.Sequential([
  tf.keras.layers.Dense(16, activation='relu'),
  tf.keras.layers.Dense(3, activation='sigmoid')
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

medium

A. Last layer activation should be softmax, not sigmoid

B. Loss function should be binary_crossentropy

C. Optimizer 'adam' is invalid

D. Dense layer units must be 1 for multi-class

Solution

1. Step 1: Check last layer activation for multi-class

   Sigmoid outputs independent probabilities, not suitable for multi-class where classes are exclusive.

2. Step 2: Correct activation for multi-class classification

   Softmax outputs probabilities summing to 1, appropriate for multi-class classification.

3. Final Answer:

   Last layer activation should be softmax, not sigmoid -> Option A

4. Quick Check:

   Multi-class needs softmax activation [OK]

Hint: Use softmax activation for multi-class last layer [OK]

Common Mistakes:

• Using sigmoid activation for multi-class output
• Confusing loss functions for classification types
• Thinking optimizer name 'adam' is invalid

5.

You want to build a multi-class classification model with 5 classes. Your labels are integers from 0 to 4. Which of the following code snippets correctly defines and compiles the model?

Option A:
model = tf.keras.Sequential([
  tf.keras.layers.Dense(32, activation='relu'),
  tf.keras.layers.Dense(5, activation='softmax')
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

Option B:
model = tf.keras.Sequential([
  tf.keras.layers.Dense(32, activation='relu'),
  tf.keras.layers.Dense(5, activation='sigmoid')
])
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

Option C:
model = tf.keras.Sequential([
  tf.keras.layers.Dense(32, activation='relu'),
  tf.keras.layers.Dense(1, activation='softmax')
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

Option D:
model = tf.keras.Sequential([
  tf.keras.layers.Dense(32, activation='relu'),
  tf.keras.layers.Dense(5, activation='softmax')
])
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

hard

A. Option A

B. Option B

C. Option C

D. Option D

Solution

1. Step 1: Check output layer units and activation

   For 5 classes, output units must be 5 with softmax activation to get class probabilities.

2. Step 2: Check loss function matches label format

   Labels are integers, so sparse_categorical_crossentropy is correct loss.

3. Step 3: Verify optimizer and metrics

   Adam optimizer and accuracy metric are appropriate choices.

4. Final Answer:

   Option A -> Option A

5. Quick Check:

   Correct output units, activation, and loss for integer labels [OK]

Hint: Match output units and loss to label format [OK]

Common Mistakes:

• Using sigmoid activation for multi-class output
• Using binary_crossentropy loss for multi-class
• Setting output units to 1 instead of number of classes