Fine-tuning approach in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Fine-tuning approach

Problem:You want to classify images into 5 categories using a pre-trained model. Currently, you trained only the top layers, but the validation accuracy is low compared to training accuracy.

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

Issue:The model is overfitting. It learns training data well but does not generalize to validation data.

Your Task

Reduce overfitting by fine-tuning some layers of the pre-trained model to improve validation accuracy to at least 80%, while keeping training accuracy below 90%.

You can only unfreeze and train layers of the pre-trained base model.

Do not change the dataset or batch size.

Keep the same optimizer and learning rate for fine-tuning.

Hint 1

Hint 2

Hint 3

Solution

TensorFlow

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

# Load pre-trained base model without top layers
base_model = MobileNetV2(input_shape=(128, 128, 3), include_top=False, weights='imagenet')
base_model.trainable = False

# Add new classification head
inputs = tf.keras.Input(shape=(128, 128, 3))
x = base_model(inputs, training=False)
x = layers.GlobalAveragePooling2D()(x)
x = layers.Dropout(0.3)(x)  # Add dropout to reduce overfitting
outputs = layers.Dense(5, activation='softmax')(x)
model = models.Model(inputs, outputs)

# Compile and train only top layers first
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
# Assume train_ds and val_ds are predefined datasets
# model.fit(train_ds, epochs=5, validation_data=val_ds)

# Fine-tuning: unfreeze last 20 layers of base model
base_model.trainable = True
for layer in base_model.layers[:-20]:
    layer.trainable = False

# Compile with lower learning rate
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-5),
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# Train with fine-tuning
# model.fit(train_ds, epochs=10, validation_data=val_ds)

# Note: Replace train_ds and val_ds with actual datasets when running.

Unfroze the last 20 layers of the pre-trained MobileNetV2 model to allow fine-tuning.

Added dropout layer with rate 0.3 before the classification layer to reduce overfitting.

Lowered the learning rate to 1e-5 during fine-tuning to avoid large weight updates.

Results Interpretation

Before fine-tuning: Training accuracy was 95%, validation accuracy was 70%, showing overfitting.

After fine-tuning: Training accuracy decreased to 88%, validation accuracy improved to 82%, and validation loss decreased, indicating better generalization.

Fine-tuning some layers of a pre-trained model with a lower learning rate and dropout helps reduce overfitting and improves validation accuracy.

Bonus Experiment

Try adding data augmentation to the training pipeline along with fine-tuning to see if validation accuracy improves further.

💡 Hint

Use layers like RandomFlip and RandomRotation from tensorflow.keras.layers to augment images during training.