Learning rate for fine-tuning in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Learning rate for fine-tuning

Problem:You want to fine-tune a pretrained image classification model on a new dataset. The current learning rate is too high, causing the model to not improve validation accuracy well.

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

Issue:The model overfits quickly and validation accuracy is low compared to training accuracy, indicating the learning rate is too high for fine-tuning.

Your Task

Adjust the learning rate to improve validation accuracy to above 80% while keeping training accuracy below 90% to reduce overfitting.

Keep the model architecture the same.

Only change the learning rate and training epochs.

Use TensorFlow and Keras API.

Hint 1

Hint 2

Hint 3

Solution

TensorFlow

import tensorflow as tf
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam

# Load pretrained base model
base_model = MobileNetV2(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
base_model.trainable = True  # Fine-tune all layers

# Add classification head
x = base_model.output
x = GlobalAveragePooling2D()(x)
outputs = Dense(10, activation='softmax')(x)  # Assume 10 classes
model = Model(inputs=base_model.input, outputs=outputs)

# Compile with a lower learning rate
optimizer = Adam(learning_rate=0.0001)  # Lowered from 0.001
model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# Dummy dataset for example (replace with real data)
import numpy as np
X_train = np.random.rand(1000, 224, 224, 3).astype('float32')
y_train = np.random.randint(0, 10, 1000)
X_val = np.random.rand(200, 224, 224, 3).astype('float32')
y_val = np.random.randint(0, 10, 200)

# Train with fewer epochs
history = model.fit(X_train, y_train, epochs=10, batch_size=32, validation_data=(X_val, y_val))

Reduced learning rate from 0.001 to 0.0001 for fine-tuning.

Reduced number of training epochs from 20 to 10 to prevent overfitting.

Results Interpretation

Before: Training accuracy 95%, Validation accuracy 70%, Training loss 0.15, Validation loss 0.85

After: Training accuracy 88%, Validation accuracy 82%, Training loss 0.30, Validation loss 0.45

Lowering the learning rate during fine-tuning helps the model adjust weights more gently, reducing overfitting and improving validation accuracy.

Bonus Experiment

Try using a learning rate scheduler that decreases the learning rate during training to see if validation accuracy improves further.

💡 Hint

Use TensorFlow's LearningRateScheduler callback or ReduceLROnPlateau to adjust learning rate dynamically.