How can Tensorflow be used to attach a classification head using Python?

TensorFlow can be used to attach a classification head using a sequential model that contains a Dense layer and a pre-defined feature extractor model. This process is essential in transfer learning where we leverage pre-trained models and add custom classification layers.

Read More: What is TensorFlow and how Keras work with TensorFlow to create Neural Networks?

What is Transfer Learning?

Transfer learning allows us to use pre-trained models as feature extractors. A model trained on a large dataset like ImageNet has already learned useful feature representations, so we don't need to train from scratch.

TensorFlow Hub provides a repository of pre-trained models that can be easily integrated with tf.keras. TensorFlow can be used to fine-tune learning models.

Attaching a Classification Head

Here's how to attach a classification head to a pre-trained feature extractor ?

import tensorflow as tf
import tensorflow_hub as hub

# Load a pre-trained feature extractor from TensorFlow Hub
feature_extractor_url = "https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4"
feature_extractor_layer = hub.KerasLayer(feature_extractor_url,
                                         input_shape=(224, 224, 3),
                                         trainable=False)

# Define class names for demonstration
class_names = ['cat', 'dog', 'bird', 'fish', 'horse']

print("Attaching a classification head")
num_classes = len(class_names)

model = tf.keras.Sequential([
    feature_extractor_layer,
    tf.keras.layers.Dense(num_classes)
])

print("The base architecture of the model")
model.summary()

# Create sample batch for prediction
import numpy as np
image_batch = np.random.random((32, 224, 224, 3))

print("The predictions are made")
predictions = model(image_batch)
print("The dimensions of the predictions")
print(predictions.shape)

The output shows the model architecture and prediction dimensions ?

Attaching a classification head
The base architecture of the model
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
keras_layer (KerasLayer)     (None, 1280)              2257984   
_________________________________________________________________
dense (Dense)                (None, 5)                 6405      
=================================================================
Total params: 2,264,389
Trainable params: 6,405
Non-trainable params: 2,257,984
_________________________________________________________________
The predictions are made
The dimensions of the predictions
(32, 5)

How It Works

The architecture consists of two main components:

• Feature Extractor Layer: Pre-trained model that converts images into 1280-dimensional feature vectors
• Dense Classification Head: Fully connected layer that maps features to class probabilities
• Trainable Parameters: Only the Dense layer (6,405 params) will be trained, while the feature extractor remains frozen
• Output Shape: (32, 5) indicates 32 samples with 5 class predictions each

Key Benefits

This approach offers several advantages:

• Faster Training: Only the classification head needs training
• Less Data Required: Pre-trained features work well with smaller datasets
• Better Performance: Leverages knowledge from large-scale pre-training

Conclusion

Attaching a classification head to a pre-trained feature extractor is an efficient way to build image classifiers. The feature extractor provides rich representations while the Dense layer learns task-specific classifications with minimal training.