Deploying 🤗 ViT on Vertex AI
Sayak Paul, chansung park
In the previous posts, we showed how to deploy a Vision Transformers(ViT) modelfrom 🤗 Transformers locally andon a Kubernetes cluster. This post willshow you how to deploy the same model on the Vertex AI platform.You’ll achieve the same scalability level as Kubernetes-based deployment but withsignificantly less code. This post builds on top of the previous two posts linked above. You’readvised to check them out if you haven’t already.You can find a completely worked-out example in the Colab Notebooklinked at the beginning of the post.What is Vertex AI?According to Google Cloud:Vertex AI provides tools to support your entire ML workflow, acrossdifferent model types and varying levels of ML expertise.Concerning model deployment, Vertex AI provides a few important featureswith a unified API design:AuthenticationAutoscaling based on trafficModel versioningTraffic splitting between different versions of a modelRate limitingModel monitoring and loggingSupport for online and batch predictionsFor TensorFlow models, it offers various off-the-shelf utilities, whichyou’ll get to in this post. But it also has similar support for otherframeworks likePyTorchand scikit-learn.To use Vertex AI, you’ll need a billing-enabled Google CloudPlatform (GCP) projectand the following services enabled:Vertex AICloud StorageRevisiting the Serving ModelYou’ll use the same ViT B/16 model implemented in TensorFlow as you did in the last two posts. You serialized the model withcorresponding pre-processing and post-processing operations embedded toreduce training-serving skew.Please refer to the first post that discussesthis in detail. The signature of the final serialized SavedModel looks like:The given SavedModel SignatureDef contains the following input(s):inputs['string_input'] tensor_info:dtype: DT_STRINGshape: (-1)name: serving_default_string_input:0The given SavedModel SignatureDef contains the following output(s):outputs['confidence'] tensor_info:dtype: DT_FLOATshape: (-1)name: StatefulPartitionedCall:0outputs['label'] tensor_info:dtype: DT_STRINGshape: (-1)name: StatefulPartitionedCall:1Method name is: tensorflow/serving/predictThe model will accept base64 encoded strings of images, performpre-processing, run inference, and finally perform the post-processingsteps. The strings are base64 encoded to prevent anymodifications during network transmission. Pre-processing includesresizing the input image to 224x224 resolution, standardizing it to the[-1, 1] range, and transposing it to the channels_first memorylayout. Postprocessing includes mapping the predicted logits to stringlabels.To perform a deployment on Vertex AI, you need to keep the modelartifacts in a Google Cloud Storage (GCS) bucket.The accompanying Colab Notebook shows how to create a GCS bucket andsave the model artifacts into it.Deployment workflow with Vertex AIThe figure below gives a pictorial workflow of deploying an alreadytrained TensorFlow model on Vertex AI.Let’s now discuss what the Vertex AI Model Registry and Endpoint are.Vertex AI Model RegistryVertex AI Model Registry is a fully managed machine learning modelregistry. There are a couple of things to note about what fully managedmeans here. First, you don’t need to worry about how and where modelsare stored. Second, it manages different versions of the same model.These features are important for machine learning in production.Building a model registry that guarantees high availability and securityis nontrivial. Also, there are often situations where you want to rollback the current model to a past version since we can notcontrol the inside of a black box machine learning model. Vertex AIModel Registry allows us to achieve these without much difficulty.The currently supported model types include SavedModel fromTensorFlow, scikit-learn, and XGBoost.Vertex AI EndpointFrom the user’s perspective, Vertex AI Endpoint simply provides anendpoint to receive requests and send responses back. However, it has alot of things under the hood for machine learning operators toconfigure. Here are some of the configurations that you can choose:Version of a modelSpecification of VM in terms of CPU, memory, and acceleratorsMin/Max number of compute nodesTraffic split percentageModel monitoring window length and its objectivesPrediction requests sampling ratePerforming the DeploymentThe google-cloud-aiplatformPython SDK provides easy APIs to manage the lifecycle of a deployment onVertex AI. It is divided into four steps:uploading a modelcreating an endpoint deploying the model to the endpointmaking prediction requests.Throughout these steps, you willneed ModelServiceClient, EndpointServiceClient, andPredictionServiceClient modules from the google-cloud-aiplatformPython SDK to interact with Vertex AI.1. The first step in the workflow is to upload the SavedModel toVertex AI’s model registry:tf28_gpu_model_dict = {"display_name": "ViT Base TF2.8 GPU model","artifact_uri": f"{GCS_BUCKET}/{LOCAL_MODEL_DIR}","container_spec": {"image_uri": "us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-8:latest",},}tf28_gpu_model = (model_service_client.upload_model(parent=PARENT, model=tf28_gpu_model_dict).result(timeout=180).model)Let’s unpack the code piece by piece:GCS_BUCKET denotes the path of your GCS bucket where the modelartifacts are located (e.g., gs://hf-tf-vision).In container_spec, you provide the URI of a Docker image that willbe used to serve predictions. Vertex AIprovides pre-built images to serve TensorFlow models, but you can also use your custom Docker images when using a different framework (an example).model_service_client is aModelServiceClientobject that exposes the methods to upload a model to the Vertex AIModel Registry.PARENT is set to f"projects/{PROJECT_ID}/locations/{REGION}"that lets Vertex AI determine where the model is going to be scopedinside GCP.2. Then you need to create a Vertex AI Endpoint:tf28_gpu_endpoint_dict = {"display_name": "ViT Base TF2.8 GPU endpoint",}tf28_gpu_endpoint = (endpoint_service_client.create_endpoint(parent=PARENT, endpoint=tf28_gpu_endpoint_dict).result(timeout=300).name)Here you’re using an endpoint_service_client which is anEndpointServiceClientobject. It lets you create and configure your Vertex AI Endpoint.3. Now you’re down to performing the actual deployment!tf28_gpu_deployed_model_dict = {"model": tf28_gpu_model,"display_name": "ViT Base TF2.8 GPU deployed model","dedicated_resources": {"min_replica_count": 1,"max_replica_count": 1,"machine_spec": {"machine_type": DEPLOY_COMPUTE, # "n1-standard-8""accelerator_type": DEPLOY_GPU, # aip.AcceleratorType.NVIDIA_TESLA_T4"accelerator_count": 1,},},}tf28_gpu_deployed_model = endpoint_service_client.deploy_model(endpoint=tf28_gpu_endpoint,deployed_model=tf28_gpu_deployed_model_dict,traffic_split={"0": 100},).result()Here, you’re chaining together the model you uploaded to the Vertex AIModel Registry and the Endpoint you created in the above steps. You’refirst defining the configurations of the deployment undertf28_gpu_deployed_model_dict.Under dedicated_resources you’re configuring:min_replica_count and max_replica_count that handle theautoscaling aspects of your deployment.machine_spec lets you define the configurations of the deploymenthardware:machine_type is the base machine type that will be used to runthe Docker image. The underlying autoscaler will scale thismachine as per the traffic load. You can choose one from the supported machine types.accelerator_type is the hardware accelerator that will be usedto perform inference.accelerator_count denotes the number of hardware accelerators toattach to each replica.Note that providing an accelerator is not a requirement to deploymodels on Vertex AI.Next, you deploy the endpoint using the above specifications:tf28_gpu_deployed_model = endpoint_service_client.deploy_model(endpoint=tf28_gpu_endpoint,deployed_model=tf28_gpu_deployed_model_dict,traffic_split={"0": 100},).result()Notice how you’re defining the traffic split for the model. If you hadmultiple versions of the model, you could have defined a dictionarywhere the keys would denote the model version and values would denotethe percentage of traffic the model is supposed to serve.With a Model Registry and a dedicatedinterfaceto manage Endpoints, Vertex AI lets you easily control the importantaspects of the deployment.It takes about 15 - 30 minutes for Vertex AI to scope the deployment.Once it’s done, you should be able to see it on theconsole.Performing PredictionsIf your deployment was successful, you can test the deployedEndpoint by making a prediction request.First, prepare a base64 encoded image string:import base64import tensorflow as tfimage_path = tf.keras.utils.get_file("image.jpg", "http://images.cocodataset.org/val2017/000000039769.jpg")bytes = tf.io.read_file(image_path)b64str = base64.b64encode(bytes.numpy()).decode("utf-8")4. The following utility first prepares a list of instances (onlyone instance in this case) and then uses a prediction service client (oftype PredictionServiceClient).serving_input is the name of the input signature key of the servedmodel. In this case, the serving_input is string_input, whichyou can verify from the SavedModel signature output shown above.from google.protobuf import json_formatfrom google.protobuf.struct_pb2 import Valuedef predict_image(image, endpoint, serving_input):# The format of each instance should conform to# the deployed model's prediction input schema.instances_list = [{serving_input: {"b64": image}}]instances = [json_format.ParseDict(s, Value()) for s in instances_list]print(prediction_service_client.predict(endpoint=endpoint,instances=instances,))predict_image(b64str, tf28_gpu_endpoint, serving_input)For TensorFlow models deployed on Vertex AI, the request payload needsto be formatted in a certain way. For models like ViT that deal withbinary data like images, they need to be base64 encoded. According tothe official guide,the request payload for each instance needs to be like so:{serving_input: {"b64": base64.b64encode(jpeg_data).decode()}}The predict_image() utility prepares the request payload conformingto this specification.If everything goes well with the deployment, when you callpredict_image(), you should get an output like so:predictions {struct_value {fields {key: "confidence"value {number_value: 0.896659553}}fields {key: "label"value {string_value: "Egyptian cat"}}}}deployed_model_id: "5163311002082607104"model: "projects/29880397572/locations/us-central1/models/7235960789184544768"model_display_name: "ViT Base TF2.8 GPU model"Note, however, this is not the only way to obtain predictions using aVertex AI Endpoint. If you head over to the Endpoint console and selectyour endpoint, it will show you two different ways to obtainpredictions:It’s also possible to avoid cURL requests and obtain predictionsprogrammatically without using the Vertex AI SDK. Refer tothis notebookto learn more.Now that you’ve learned how to use Vertex AI to deploy a TensorFlowmodel, let’s now discuss some beneficial features provided by Vertex AI.These help you get deeper insights into your deployment.Monitoring with Vertex AIVertex AI also lets you monitor your model without any configuration.From the Endpoint console, you can get details about the performance ofthe Endpoint and the utilization of the allocated resources.As seen in the above chart, for a brief amount of time, the acceleratorduty cycle (utilization) was about 100% which is a sight for sore eyes.For the rest of the time, there weren’t any requests to process hencethings were idle.This type of monitoring helps you quickly flag the currently deployedEndpoint and make adjustments as necessary. It’s also possible torequest monitoring of model explanations. Referhereto learn more.Local Load TestingWe conducted a local load test to better understand the limits of theEndpoint with Locust. The table belowsummarizes the request statistics:Among all the different statistics shown in the table, Average (ms)refers to the average latency of the Endpoint. Locust fired off about17230 requests, and the reported average latency is 646Milliseconds, which is impressive. In practice, you’d want to simulatemore real traffic by conducting the load test in a distributed manner.Refer hereto learn more.This directoryhas all the information needed to know how we conducted the load test.PricingYou can use the GCP cost estimator to estimate the cost of usage, and the exact hourly pricing table can be found here.It is worth noting that you are only charged when the node is processingthe actual prediction requests, and you need to calculate the price withand without GPUs.For the Vertex Prediction for a custom-trained model, we can chooseN1 machine types from n1-standard-2 to n1-highcpu-32.You used n1-standard-8 for this post which is equipped with 8vCPUs and 32GBs of RAM.Machine TypeHourly Pricing (USD)n1-standard-8 (8vCPU, 30GB)$ 0.4372Also, when you attach accelerators to the compute node, you will becharged extra by the type of accelerator you want. We usedNVIDIA_TESLA_T4 for this blog post, but almost all modernaccelerators, including TPUs are supported. You can find furtherinformation here.Accelerator TypeHourly Pricing (USD)NVIDIA_TESLA_T4$ 0.4024Call for ActionThe collection of TensorFlow vision models in 🤗 Transformers is growing. It now supportsstate-of-the-art semantic segmentation with SegFormer.We encourage you to extend the deployment workflow you learned in this post to semantic segmentation models like SegFormer.ConclusionIn this post, you learned how to deploy a Vision Transformer model withthe Vertex AI platform using the easy APIs it provides. You also learnedhow Vertex AI’s features benefit the model deployment process by enablingyou to focus on declarative configurations and removing the complexparts. Vertex AI also supports deployment of PyTorch models via customprediction routes. Referherefor more details.The series first introduced you to TensorFlow Serving for locally deploying a vision model from 🤗 Transformers. In the second post, you learned how to scalethat local deployment with Docker and Kubernetes. We hope this series on theonline deployment of TensorFlow vision models was beneficial for you to take yourML toolbox to the next level. We can’t wait to see what you build with these tools.AcknowledgementsThanks to the ML Developer Relations Program team at Google, whichprovided us with GCP credits for conducting the experiments.Parts of the deployment code were referred from this notebook of the official GitHub repositoryof Vertex AI code samples.