Deploying TensorFlow Vision Models in Hugging Face with TF Serving

In the past few months, the Hugging Face team and external contributors added a variety of vision models in TensorFlow to Transformers. This list is growing comprehensively and already includes state-of-the-art pre-trained models like Vision Transformer, Masked Autoencoders, RegNet, ConvNeXt, and many others!

When it comes to deploying TensorFlow models, you have got a variety of options. Depending on your use case, you may want to expose your model as an endpoint or package it in an application itself. TensorFlow provides tools that cater to each of these different scenarios.

In this post, you'll see how to deploy a Vision Transformer (ViT) model (for image classification) locally using TensorFlow Serving (TF Serving). This will allow developers to expose the model either as a REST or gRPC endpoint. Moreover, TF Serving supports many deployment-specific features off-the-shelf such as model warmup, server-side batching, etc.

To get the complete working code shown throughout this post, refer to the Colab Notebook shown at the beginning.

Saving the Model

All TensorFlow models in 🤗 Transformers have a method named save_pretrained(). With it, you can serialize the model weights in the h5 format as well as in the standalone SavedModel format. TF Serving needs a model to be present in the SavedModel format. So, let's first load a Vision Transformer model and save it:

from transformers import TFViTForImageClassification

temp_model_dir = "vit"
ckpt = "google/vit-base-patch16-224"

model = TFViTForImageClassification.from_pretrained(ckpt)
model.save_pretrained(temp_model_dir, saved_model=True)

By default, save_pretrained() will first create a version directory inside the path we provide to it. So, the path ultimately becomes: {temp_model_dir}/saved_model/{version}.

We can inspect the serving signature of the SavedModel like so:

saved_model_cli show --dir {temp_model_dir}/saved_model/1 --tag_set serve --signature_def serving_default

This should output:

Model Surgery

This should print: