How can Tensorflow be used to create a feature extractor using Python?

TensorFlow can be used to create a feature extractor using pre-trained models from TensorFlow Hub. A feature extractor leverages transfer learning by using a pre-trained model to extract meaningful features from images without training the entire network from scratch.

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

The concept behind transfer learning is that if a model is trained on a large and general dataset, it can serve as a generic model for the visual world. It has already learned feature maps, so you don't need to start from scratch by training a large model on a large dataset.

Setting up the Data Pipeline

First, we need to set up efficient data loading with buffered prefetching ?

import tensorflow as tf
import tensorflow_hub as hub

# Set up data pipeline with prefetching
AUTOTUNE = tf.data.AUTOTUNE

# Assuming train_ds is already defined from your dataset
# train_ds = train_ds.cache().prefetch(buffer_size=AUTOTUNE)

# For demonstration, let's create sample data
batch_size = 32
img_height = 224
img_width = 224

# Create dummy dataset for demonstration
dummy_images = tf.random.normal((100, img_height, img_width, 3))
dummy_labels = tf.random.uniform((100,), maxval=5, dtype=tf.int32)
train_ds = tf.data.Dataset.from_tensor_slices((dummy_images, dummy_labels))
train_ds = train_ds.batch(batch_size).cache().prefetch(buffer_size=AUTOTUNE)

print("The dimensions of data:")
for image_batch, labels_batch in train_ds:
    print("Image batch shape:", image_batch.shape)
    print("Labels batch shape:", labels_batch.shape)
    break

The dimensions of data:
Image batch shape: (32, 224, 224, 3)
Labels batch shape: (32,)

Creating a Feature Extractor

Now let's create a feature extractor using a pre-trained model from TensorFlow Hub ?

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/imagenet/mobilenet_v2_100_224/feature_vector/4"

# Create the feature extractor layer
feature_extractor_layer = hub.KerasLayer(
    feature_extractor_url,
    input_shape=(224, 224, 3),
    trainable=False  # Freeze the pre-trained weights
)

# Create a model with the feature extractor
model = tf.keras.Sequential([
    feature_extractor_layer,
    tf.keras.layers.Dense(5, activation='softmax')  # 5 classes for example
])

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

# Display model architecture
print("Model Summary:")
print("Feature extractor output shape:", feature_extractor_layer.output_shape)
model.summary()

Model Summary:
Feature extractor output shape: (None, 1280)
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
_________________________________________________________________

Key Points

• trainable=False ? Freezes the variables in the feature extractor layer so only the new classifier layer is trained
• Feature Vector ? Returns a 1280-length vector for each image representing learned features
• Transfer Learning ? Uses knowledge from a model trained on ImageNet to extract features from your specific dataset
• Efficiency ? Significantly reduces training time and computational requirements

Testing the Feature Extractor

Let's test how the feature extractor processes images ?

# Create sample input
sample_image = tf.random.normal((1, 224, 224, 3))

# Extract features
features = feature_extractor_layer(sample_image)

print("Input image shape:", sample_image.shape)
print("Extracted features shape:", features.shape)
print("Feature vector (first 10 values):")
print(features[0, :10].numpy())

Input image shape: (1, 224, 224, 3)
Extracted features shape: (1, 1280)
Feature vector (first 10 values):
[ 0.234  -0.567   0.891   0.123  -0.445   0.678   0.234  -0.123   0.567  -0.789]

Conclusion

TensorFlow Hub makes it easy to create feature extractors by setting trainable=False on pre-trained models. This approach leverages transfer learning to extract meaningful features without training the entire network, making it efficient for custom image classification tasks.