在Vertex AI上部署🤗 ViT

Deploy ViT on Vertex AI

在之前的帖子中,我们展示了如何从🤗 Transformers本地和在Kubernetes集群上部署Vision Transformers(ViT)模型。本文将向您展示如何在Vertex AI平台上部署相同的模型。您将实现与基于Kubernetes的部署相同的可扩展性水平,但代码量要少得多。

本文基于上面链接的前两篇帖子。如果您还没有查看它们,建议您先查看一下。

您可以在本文开头链接的Colab Notebook中找到一个完整的示例。

根据Google Cloud的说法:

Vertex AI提供了支持整个ML工作流程的工具,包括不同模型类型和不同ML专业水平。

关于模型部署,Vertex AI提供了几个重要的功能,并具有统一的API设计:

  • 身份验证

  • 基于流量的自动扩展

  • 模型版本控制

  • 在不同模型版本之间进行流量分割

  • 限制速率

  • 模型监控和日志记录

  • 支持在线和批量预测

对于TensorFlow模型,它提供了各种现成的实用工具,您将在本文中了解到。但是它也对其他框架(如PyTorch和scikit-learn)提供类似的支持。

要使用Vertex AI,您需要启用计费的Google Cloud Platform(GCP)项目和以下服务:

  • Vertex AI

  • Cloud Storage

您将使用前两篇帖子中使用的相同的ViT B/16模型,该模型在TensorFlow中实现。您已经将模型与相应的预处理和后处理操作一起序列化,以减少训练-服务偏差。请参阅详细讨论此问题的第一篇帖子。最终序列化的SavedModel的签名如下所示:

给定的SavedModel SignatureDef包含以下输入:
  inputs['string_input'] tensor_info:
      dtype: DT_STRING
      shape: (-1)
      name: serving_default_string_input:0
给定的SavedModel SignatureDef包含以下输出:
  outputs['confidence'] tensor_info:
      dtype: DT_FLOAT
      shape: (-1)
      name: StatefulPartitionedCall:0
  outputs['label'] tensor_info:
      dtype: DT_STRING
      shape: (-1)
      name: StatefulPartitionedCall:1
方法名称是:tensorflow/serving/predict

该模型将接受图像的base64编码字符串,进行预处理,运行推断,最后执行后处理步骤。字符串进行base64编码以防止在网络传输过程中进行任何修改。预处理包括将输入图像调整为224×224分辨率,将其标准化为[-1, 1]范围,并将其转置为channels_first内存布局。后处理包括将预测的逻辑映射到字符串标签。

要在Vertex AI上执行部署,您需要将模型工件保留在Google Cloud Storage(GCS)存储桶中。随附的Colab Notebook显示了如何创建GCS存储桶并将模型工件保存到其中。

下图给出了在Vertex AI上部署已经训练好的TensorFlow模型的流程图。

现在让我们讨论一下Vertex AI模型注册表和端点。

Vertex AI模型注册表

Vertex AI模型注册表是一个完全托管的机器学习模型注册表。关于这里完全托管的几点需要注意。首先,您不需要担心模型的存储方式和位置。其次,它管理同一模型的不同版本。

这些功能对于生产中的机器学习非常重要。构建一个能够保证高可用性和安全性的模型注册表并不是一件简单的事情。此外,由于我们无法控制黑盒机器学习模型的内部情况,通常会出现想要将当前模型回滚到过去版本的情况。Vertex AI模型注册表使我们能够在没有太多困难的情况下实现这些目标。

目前支持的模型类型包括来自TensorFlow、scikit-learn和XGBoost的SavedModel

Vertex AI 终端

从用户的角度来看,Vertex AI 终端只是提供一个接收请求并发送响应的终端。然而,在机器学习操作员的配置下,它在幕后做了很多事情。以下是一些可供您选择的配置:

  • 模型的版本

  • VM 的规格,包括 CPU、内存和加速器

  • 计算节点的最小/最大数量

  • 流量分配百分比

  • 模型监控窗口的长度和目标

  • 预测请求的采样率

google-cloud-aiplatform Python SDK 提供了简单的 API 来管理 Vertex AI 上部署的生命周期。它分为四个步骤:

  1. 上传模型
  2. 创建终端
  3. 将模型部署到终端
  4. 进行预测请求。

在这些步骤中,您需要使用 google-cloud-aiplatform Python SDK 中的 ModelServiceClientEndpointServiceClientPredictionServiceClient 模块与 Vertex AI 进行交互。

1. 工作流程的第一步是将 SavedModel 上传到 Vertex AI 的模型注册表中:

tf28_gpu_model_dict = {
    "display_name": "ViT Base TF2.8 GPU 模型",
    "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
)

让我们逐段解读代码:

  • GCS_BUCKET 表示存放模型工件的 GCS 存储桶的路径(例如:gs://hf-tf-vision)。

  • container_spec 中,您提供将用于提供预测的 Docker 镜像的 URI。Vertex AI 提供了预先构建的图像来提供 TensorFlow 模型的服务,但是当使用不同的框架时,您也可以使用自定义的 Docker 镜像(示例)。

  • model_service_client 是一个 ModelServiceClient 对象,它提供了将模型上传到 Vertex AI 模型注册表的方法。

  • PARENT 设置为 f"projects/{PROJECT_ID}/locations/{REGION}",让 Vertex AI 确定模型将在 GCP 中的范围。

2. 然后您需要创建一个 Vertex AI 终端:

tf28_gpu_endpoint_dict = {
    "display_name": "ViT Base TF2.8 GPU 终端",
}
tf28_gpu_endpoint = (
    endpoint_service_client.create_endpoint(
        parent=PARENT, endpoint=tf28_gpu_endpoint_dict
    )
    .result(timeout=300)
    .name
)

在这里,您使用的是一个 endpoint_service_client,它是一个 EndpointServiceClient 对象。它允许您创建和配置 Vertex AI 终端。

3. 现在您可以进行实际部署了!

tf28_gpu_deployed_model_dict = {
    "model": tf28_gpu_model,
    "display_name": "ViT Base TF2.8 GPU 部署的模型",
    "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()

在这里,您将上传到Vertex AI模型注册表和上述步骤中创建的端点连接在一起。首先,在tf28_gpu_deployed_model_dict下定义部署的配置。

dedicated_resources下进行配置:

  • min_replica_countmax_replica_count处理部署的自动缩放方面。

  • machine_spec允许您定义部署硬件的配置:

    • machine_type是用于运行Docker镜像的基本机器类型。基础自动缩放器将根据流量负载扩展此机器。您可以从支持的机器类型中选择一个。

    • accelerator_type是用于执行推断的硬件加速器。

    • accelerator_count表示每个副本附加的硬件加速器数量。

注意,在Vertex AI上部署模型时提供加速器不是必需的。

接下来,使用上述规格部署端点:

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()

请注意,您正在为模型定义流量分割。如果您有多个模型版本,您可以定义一个字典,其中键表示模型版本,值表示模型应该提供的流量百分比。

通过模型注册表和专用的端点管理界面,Vertex AI可以轻松控制部署的重要方面。

Vertex AI需要大约15-30分钟来为部署提供范围。完成后,您应该能够在控制台上看到它。

如果部署成功,您可以通过进行预测请求来测试部署的端点。

首先,准备一个base64编码的图像字符串:

import base64
import tensorflow as tf

image_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.以下实用程序首先准备一个实例列表(在本例中仅有一个实例),然后使用一个预测服务客户端(类型为PredictionServiceClient)。serving_input是提供的模型输入签名键的名称。在本例中,serving_inputstring_input,您可以从上面显示的SavedModel签名输出中验证。

from google.protobuf import json_format
from google.protobuf.struct_pb2 import Value

def predict_image(image, endpoint, serving_input):
    # 每个实例的格式应符合部署模型的预测输入模式。
    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)

对于在Vertex AI上部署的TensorFlow模型,请求负载需要以特定的方式进行格式化。对于像ViT这样处理像图像这样的二进制数据的模型,它们需要进行base64编码。根据官方指南,每个实例的请求负载应该是这样的:

{serving_input: {"b64": base64.b64encode(jpeg_data).decode()}}

predict_image()实用程序准备符合此规范的请求负载。

如果部署一切顺利,当您调用predict_image()时,您应该会得到以下输出:

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"

请注意,这不是使用Vertex AI Endpoint获取预测的唯一方法。如果您转到Endpoint控制台并选择您的端点,它将向您展示获取预测的两种不同方法:

还可以避免使用cURL请求并以编程方式获取预测,而无需使用Vertex AI SDK。请参考这个笔记本以了解更多信息。

现在您已经学会如何使用Vertex AI部署TensorFlow模型,让我们讨论一些Vertex AI提供的有益功能。这些功能可以帮助您更深入地了解您的部署。

Vertex AI还允许您在不进行任何配置的情况下监视您的模型。您可以从Endpoint控制台获取有关端点性能和分配资源利用情况的详细信息。

从上图可以看出,在短暂的时间内,加速器工作周期(利用率)达到了约100%,这是令人欣慰的。其余时间内没有任何请求需要处理,因此处于空闲状态。

这种类型的监视帮助您快速标记当前部署的端点,并根据需要进行调整。还可以请求对模型解释进行监视。请参考此处以了解更多信息。

我们进行了本地负载测试,以更好地了解使用Locust的端点限制。下表总结了请求统计信息:

在表中显示的各种统计数据中,Average (ms)指的是端点的平均延迟。Locust发出了约17230个请求,报告的平均延迟为646毫秒,令人印象深刻。在实践中,您希望通过分布式方式进行负载测试,以模拟更多真实的流量。请参考此处以了解更多信息。

此目录包含了了解我们如何进行负载测试所需的所有信息。

您可以使用GCP成本估算器估算使用的成本,具体的小时定价表可以在此处找到。值得注意的是,只有在节点处理实际预测请求时才会收费,并且您需要计算带有和不带有GPU的价格。

对于自定义训练模型的Vertex Prediction,我们可以从n1-standard-2n1-highcpu-32选择N1机器类型。在本文中,您使用了配备8个vCPU和32GB内存的n1-standard-8

此外,当您将加速器附加到计算节点时,您将按照所需的加速器类型额外收费。在本博客文章中,我们使用了NVIDIA_TESLA_T4,但几乎所有现代加速器,包括TPU,都得到了支持。您可以在此处找到更多信息。

🤗 Transformers中的TensorFlow视觉模型集合正在不断增长。它现在支持使用SegFormer进行最先进的语义分割。我们鼓励您将在本文中学到的部署工作流扩展到像SegFormer这样的语义分割模型。

在本文中,您学会了如何使用Vertex AI平台提供的简单API部署Vision Transformer模型。您还了解了Vertex AI的特性如何通过使您能够专注于声明性配置并消除复杂部分来受益于模型部署过程。Vertex AI还通过自定义预测路由支持部署PyTorch模型。请参考此处以了解更多详细信息。

该系列首先介绍了使用🤗 Transformers从本地部署视觉模型的TensorFlow Serving。在第二篇文章中,您了解了如何使用Docker和Kubernetes扩展该本地部署。我们希望这个关于在线部署TensorFlow视觉模型的系列对您将机器学习工具箱提升到新水平有所帮助。我们迫不及待地想看看您如何使用这些工具建立起来。

感谢Google的ML开发者关系项目团队为我们提供GCP积分以进行实验。

部署代码的部分内容参考了Vertex AI代码示例的官方GitHub存储库中的此笔记本。