TensorFlow Hub: Home to 100s of pre-trained models

Detailed explanation about TensorFlow Hub, Installation, and Usage with an example

Photo by Lyman Hansel Gerona on Unsplash

To access TensorFlow Hub click here — https://www.tensorflow.org/hub

Introduction

TensorFlow Hub is a library for the publication, discovery, and consumption of reusable models in TensorFlow. It provides a simple way to use pre-trained models for a variety of tasks, such as image classification, text analysis, and more.

TensorFlow Hub provides a wide range of pre-trained models developed by researchers and engineers from TensorFlow and the broader machine-learning community.

Here are some examples of the types of models available in TensorFlow Hub:

1. Image classification models: These models are trained on large datasets of labeled images and can classify images into various categories. Some popular image classification models in TensorFlow Hub include Inception, MobileNet, ResNet, and VGG.
2. Object detection models: These models can detect and localize objects within an image. Some popular object detection models in TensorFlow Hub include Faster R-CNN and YOLO.
3. Natural language processing (NLP) models: These models can analyze text and perform tasks such as sentiment analysis, text classification, and language translation. Some popular NLP models in TensorFlow Hub include BERT and ALBERT.
4. Speech recognition models: These models can transcribe speech into text. Some popular speech recognition models in TensorFlow Hub include wav2vec2, spice, and yamnet.
5. Generative models: These models can generate new content, such as images or text, based on input data. Some popular generative models in TensorFlow Hub include Progressive GAN and BigGAN.
6. Transfer learning models: These models are pre-trained on large datasets and can be fine-tuned for a specific task using a smaller dataset. Transfer learning models in TensorFlow Hub can be used for various tasks such as image classification, object detection, segmentation, and text analysis.

TensorFlow Hub provides a convenient way to access these pre-trained models and use them for various machine-learning tasks. The models are available in various formats, such as TensorFlow SavedModel, Keras models, and TensorFlow.js models, making it easy to integrate them into your machine learning pipeline.

Installation and Usage

1. Install TensorFlow and TensorFlow Hub:

Before using TensorFlow Hub, you need to install both TensorFlow and TensorFlow Hub. You can install them using pip in your command prompt or terminal:

pip install tensorflow
pip install tensorflow-hub

This will install the latest versions of TensorFlow and TensorFlow Hub.

Load a pre-trained model from TensorFlow Hub:

You can browse the available models on the TensorFlow Hub website (https://tfhub.dev/).

To load a pre-trained model, you first need to get its URL from the TensorFlow Hub website. For example, if you want to use an EfficientNet_v2_s model pre-trained on ImageNet 1K dataset for image classification, you can use the following URL:

module_url = "https://tfhub.dev/google/imagenet/efficientnet_v2_imagenet1k_s/classification/2"

Next, you can create a Keras layer using the hub.KerasLayer class and the URL:

import tensorflow as tf
import tensorflow_hub as hub

feature_extractor = hub.KerasLayer(module_url, input_shape=(384,384,3))

In this example, we are creating a Keras layer that will be used to extract features from images. The input_shape argument specifies the shape of the input images, which in this case is 384x384 pixels with 3 color channels (RGB).

1. Use the pre-trained model to perform a task:

Once you have loaded the pre-trained model, you can use it to perform a specific task. For example, if you take a random image of a tiger, you can do the following:

import numpy as np
import PIL.Image as Image
import matplotlib.pyplot as plt

image = Image.open("image.jpg").resize((384,384))
plt.imshow(image)
plt.show()
image_array = np.array(image) / 255.0
image_batch = np.expand_dims(image_array, axis=0)

features = feature_extractor(image_batch)

Output:

In this example, we are loading an image from a file, resizing it to 384x384 pixels, displaying the image, converting it to a NumPy array, and normalizing it to values between 0 and 1. We then add an extra dimension to the array to create a batch of images (as the pre-trained model expects a batch of images as input). Finally, we use the pre-trained EfficientNet_v2_s model to classify the image.

labels_file = "https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt"

#download labels and creates a maps
downloaded_file = tf.keras.utils.get_file("labels.txt", origin=labels_file)

classes = []

with open(downloaded_file) as f:
  labels = f.readlines()
  classes = [l.strip() for l in labels]

This code snippet uses the TensorFlow library to download a text file containing the ImageNet dataset labels from a Google Cloud Storage bucket. The URL for the text file is stored in the labels_file variable.tf.keras.utils.get_file() the function is then used to download the file from the URL specified labels_file and saves it locally with the name "labels.txt". The get_file() function caches the file locally, so if the same URL is requested again, the locally cached file will be used instead of downloading the file again.

After the file is downloaded, the code opens it and reads its contents using open() function. The readlines() function returns a list of strings, with each string representing a line from the file. The code then removes any leading or trailing whitespace from each line using the strip() function and stores the cleaned-up labels in a list called classes.

The resulting classes list contains the labels for the 1001 classes (1000 main classes + 1 class for ‘background’) of the ImageNet dataset. This list can be used to map the output of a neural network to the corresponding label for display or further analysis.

top_5 = tf.argsort(features, axis=-1, direction="DESCENDING")[0][:5].numpy()

features_array = np.array(features)
features_flatten = features_array.flatten()

for i, item in enumerate(top_5):
  class_index = item +1
  line = classes[class_index].upper()
  prob = round(features_flatten[item]*10,2)
  print(f'Predicted class : {line} : with probability : {prob}%')

Output:

Predicted class : TIGER : with probability : 90.79%
Predicted class : TIGER CAT : with probability : 77.36%
Predicted class : JAGUAR : with probability : 37.36%
Predicted class : LYNX : with probability : 29.64%
Predicted class : LEOPARD : with probability : 25.47%

This code snippet uses the TensorFlow library to predict the top 5 most likely ImageNet classes for an input image.

The variable features contains the output of a neural network model for an input image. The tf.argsort() function is used to sort the output values in descending order and returns the indices of the sorted values. The top 5 indices are selected by slicing the resulting tensor with [0][:5] and then converted to a NumPy array using the .numpy() method.

The features_array variable is created by converting the features tensor to a NumPy array. The features_flatten variable is created by flattening the features_array into a 1D array.

A for loop is then used to iterate over the top 5 indices. For each index, the corresponding ImageNet class label is retrieved from the classes list by adding 1 to the index (since the indices are 0-based, but the class labels are 1-based) and converting the label to uppercase.

Finally, the predicted class label and its corresponding probability are printed to the console using the print() function. The output will show the predicted class and its corresponding probability for each of the top 5 most likely ImageNet classes.

Fine-tune the pre-trained model (optional):

If the pre-trained model doesn’t fully meet your requirements, you can fine-tune it by adding new layers or retraining it on a new dataset. Here’s an example of how to fine-tune the EfficientNet_v2_s model for a new classification task:

import tensorflow as tf
import tensorflow_hub as hub
from tensorflow.keras import layers

model_url = "https://tfhub.dev/google/imagenet/efficientnet_v2_imagenet1k_s/feature_vector/2"
model = hub.KerasLayer(model_url, trainable=True)

num_classes = 10 #just an example

fine_tuning_model = tf.keras.Sequential([
    model,
    layers.Dense(num_classes, activation='softmax')
])

fine_tuning_model.compile(
    optimizer=tf.keras.optimizers.Adam(),
    loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True),
    metrics=[tf.keras.metrics.CategoricalAccuracy()]
)

epochs = 10
batch_size = 32

fine_tuning_model.fit(
    train_data,
    epochs=epochs,
    batch_size=batch_size,
    validation_data=val_data
)

test_loss, test_acc = fine_tuning_model.evaluate(test_data)
print('Test accuracy:', test_acc)

In this example, we create a new model by adding a new dense layer for classification. We then compile the model with an optimizer, a loss function, and metrics. Finally, we train the model on a new dataset for a specified number of epochs.

Conclusion

These are the basic steps for the installation and usage of TensorFlow Hub. TensorFlow Hub provides many more features, such as caching, versioning, and more, so be sure to check out the documentation for more information.