Deploying 🤗 ViT on Vertex AI

In the previous posts, we showed how to deploy a

Vision Transformers (ViT) model

from 🤗 Transformers locally and on a Kubernetes cluster. This post will show you how to deploy the same model on the

Vertex AI platform

. You’ll achieve the same scalability level as Kubernetes-based deployment but with significantly less code.

This post builds on top of the previous two posts linked above. You’re advised to check them out if you haven’t already.

You can find a completely worked-out example in the Colab Notebook linked at the beginning of the post.

What is Vertex AI?

According to Google Cloud:

Vertex AI provides tools to support your entire ML workflow, across different model types and varying levels of ML expertise.

Concerning model deployment, Vertex AI provides a few important features with a unified API design:

• Authentication

• Autoscaling based on traffic

• Model versioning

• Traffic splitting between different versions of a model

• Rate limiting

• Model monitoring and logging

• Support for online and batch predictions

For TensorFlow models, it offers various off-the-shelf utilities, which you’ll get to in this post. But it also has similar support for other frameworks like PyTorch and scikit-learn.

To use Vertex AI, you’ll need a billing-enabled Google Cloud Platform (GCP) project and the following services enabled:

• Vertex AI

• Cloud Storage

Revisiting the Serving Model

You’ll use the same ViT B/16 model implemented in TensorFlow as you did in the last two posts. You serialized the model with corresponding pre-processing and post-processing operations embedded to reduce training-serving skew. Please refer to the first post that discusses this in detail. The signature of the final serialized SavedModel looks like:

To perform a deployment on Vertex AI, you need to keep the model artifacts in a Google Cloud Storage (GCS) bucket. The accompanying Colab Notebook shows how to create a GCS bucket and save the model artifacts into it.

Deployment workflow with Vertex AI

Let’s now discuss what the Vertex AI Model Registry and Endpoint are.