Getting Started with KerasHub

Author: Matthew Watson, Jonathan Bischof Date created: 2022/12/15 Last modified: 2024/10/17 Description: An introduction to the KerasHub API.
View in Colab • GitHub source
KerasHub is a pretrained modeling library that aims to be simple, flexible, and fast. The library provides Keras 3 implementations of popular model architectures, paired with a collection of pretrained checkpoints available on Kaggle. Models can be used for both training and inference on any of the TensorFlow, Jax, and Torch backends.
KerasHub is an extension of the core Keras API; KerasHub components are provided as keras.Layers and keras.Models. If you are familiar with Keras, congratulations! You already understand most of KerasHub.
This guide is meant to be an accessible introduction to the entire library. We will start by using high-level APIs to classify images and generate text, then progressively show deeper customization of models and training. Throughout the guide, we use Professor Keras, the official Keras mascot, as a visual reference for the complexity of the material:
As always, we'll keep our Keras guides focused on real-world code examples. You can play with the code here at any time by clicking the Colab link at the top of the guide.
To begin, let's install keras-hub. The library is available on PyPI, so we can simply install it with pip.
Keras 3 was built to work on top of TensorFlow, Jax, and Torch backends. You should specify the backend first thing when writing Keras code, before any library imports. We will use the Jax backend for this guide, but you can use torch or tensorflow without changing a single line in the rest of this guide. That's the power of Keras 3!
We will also set XLA_PYTHON_CLIENT_MEM_FRACTION, which frees up the whole GPU for Jax to use from the start.
Lastly, we need to do some extra setup to access the models used in this guide. Many popular open LLMs, such as Gemma from Google and Llama from Meta, require accepting a community license before accessing the model weights. We will be using Gemma in this guide, so we can follow the following steps:
Before we begin, let's take a look at the key classes we will use in the KerasHub library.
All of the classes listed here have a from_preset() constructor, which will instantiate the component with weights and state for the given pre-trained model identifier. E.g. keras_hub.tokenizers.Tokenizer.from_preset("gemma2_2b_en") will create a layer that tokenizes text using a Gemma2 tokenizer vocabulary.
The figure below shows how all these core classes interact. Arrow indicate composition not inheritance (e.g., a task has a backbone).
Enough setup! Let's have some fun with pre-trained models. Let's load a test image of a California Quail and classify it.
We can use a ResNet vision model trained on the ImageNet-1k database. This model will give each input sample and output label from [0, 1000), where each label corresponds to some real word entity, like a "milk can" or a "porcupine." The dataset actually has a specific label for quail, at index 85. Let's download the model and predict a label.
These ImageNet labels aren't a particularly "human readable," so we can use a built-in utility function to decode the predictions to a set of class names.
Looking good! The model weights successfully downloaded, and we predicted the correct classification label for our quail image with near certainty.
This was our first example of the high-level task API mentioned in the API quickstart above. An keras_hub.models.ImageClassifier is a task for classifying images, and can be used with a number of different model architectures (ResNet, VGG, MobileNet, etc). You can view the full list of models shipped directly by the Keras team on Kaggle.
A task is just a subclass of keras.Model — you can use fit(), compile(), and save() on our classifier object same as any other model. But tasks come with a few extras provided by the KerasHub library. The first and most important is from_preset(), a special constructor you will see on many classes in KerasHub.
A preset is a directory of model state. It defines both the architecture we should load and the pretrained weights that go with it. from_preset() allows us to load preset directories from a number of different locations:
You can take a look at the keras_hub.models.ImageClassifier.from_preset docs to better understand all the options when constructing a Keras model from a preset.
All tasks use two main sub-objects. A keras_hub.models.Backbone and a keras_hub.layers.Preprocessor. You might be familiar already with the term backbone from computer vision, where it is often used to describe a feature extractor network that maps images to a latent space. A KerasHub backbone is this concept generalized, we use it to refer to any pretrained model without a task-specific head. That is, a KerasHub backbone maps raw images, audio and text (or a combination of these inputs) to a pretrained model's latent space. We can then map this latent space to any number of task specific outputs, depending on what we are trying to do with the model.
A preprocessor is just a Keras layer that does all the preprocessing for a specific task. In our case, preprocessing with will resize our input image and rescale it to the range [0, 1] using some ImageNet specific mean and variance data. Let's call our task's preprocessor and backbone in succession to see what happens to our input shape.
Our raw image is rescaled to (224, 224) during preprocessing and finally downscaled to a (7, 7) image of 2048 feature vectors — the latent space of the ResNet model. Note that ResNet can actually handle images of arbitrary sizes, though performance will eventually fall off if your image is very different sized than the pretrained data. If you'd like to disable the resizing in the preprocessing layer, you can run image_classifier.preprocessor.image_size = None.
If you are ever wondering the exact structure of the task you loaded, you can use model.summary() same as