将模型部署到 Cloud TPU 虚拟机

Google Cloud 提供对称为张量处理单元 (TPU) 的定制机器学习加速器的访问权限。TPU 经过优化,可以加快机器学习模型的训练和推理速度,非常适合包括自然语言处理、计算机视觉和语音识别在内的各种应用。

本页面介绍了如何将模型部署到单主机 Cloud TPU v5e 中,以便在 Vertex AI 中进行在线预测。

仅支持 Cloud TPU v5e。不支持其他 Cloud TPU 世代。

导入模型

如需在 Cloud TPU 上进行部署,您必须将模型导入 Vertex AI,并将其配置为使用以下容器之一:

预构建的优化 TensorFlow 运行时容器

如需在 Cloud TPU 上导入和运行 SavedModel,模型必须针对 TPU 进行优化。如果 TensorFlow SavedModel 尚未针对 TPU 进行优化,您可以通过以下三种方式优化模型:

  • 手动优化模型 - 您可以使用推理转换器来优化模型并保存它。然后,您必须在 upload 模型时传递 --saved_model_tags='serve,tpu'--disable_optimizer=true 标志。例如:

    model = aiplatform.Model.upload(
        display_name='Manually optimized model',
        artifact_uri="gs://model-artifact-uri",
        serving_container_image_uri="us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-tpu.2-15:latest",
        serving_container_args=[
            "--saved_model_tags=serve,tpu",
            "--disable_optimizer=true",
        ]
    )
    
  • 通过自动分区自动优化模型 - 导入模型时,Vertex AI 会尝试使用自动分区算法来优化未优化的模型。此优化并不适用于所有模型。如果优化失败,您必须手动优化模型,或选择通过手动分区自动优化模型。例如:

    model = aiplatform.Model.upload(
        display_name='TPU optimized model with automatic partitioning',
        artifact_uri="gs://model-artifact-uri",
        serving_container_image_uri="us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-tpu.2-15:latest",
        serving_container_args=[
        ]
    )
    
  • 通过手动分区自动优化模型。指定 --converter_options_string 标志并根据您的需求调整 ConverterOptions.TpuFunction。如需查看示例,请参阅转换器图片。请注意,仅 ConverterOptions.TpuFunction 受支持(手动分区只需此标志)。例如:

    model = aiplatform.Model.upload(
    display_name='TPU optimized model with manual partitioning',
      artifact_uri="gs://model-artifact-uri",
      serving_container_image_uri="us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-tpu.2-15:latest",
      serving_container_args=[
          "--converter_options_string='tpu_functions { function_alias: \"partitioning function name\" }'"
      ]
    )
    

如需详细了解如何导入模型,请参阅将模型导入 Vertex AI

预构建的 PyTorch 容器

在 Cloud TPU 上导入和运行 PyTorch 模型的说明与导入和运行 PyTorch 模型的说明相同。

例如,适用于 Cloud TPU v5e 推理的 TorchServe 演示了如何使用 Torch Model Archiver 将 Densenet 161 模型打包到模型工件中。

然后,将模型工件上传到 Cloud Storage 文件夹,并上传您的模型,如下所示:

model = aiplatform.Model.upload(
    display_name='DenseNet TPU model from SDK PyTorch 2.1',
    artifact_uri="gs://model-artifact-uri",
    serving_container_image_uri="us-docker.pkg.dev/vertex-ai/prediction/pytorch-tpu.2-1:latest",
    serving_container_args=[],
    serving_container_predict_route="/predictions/model",
    serving_container_health_route="/ping",
    serving_container_ports=[8080]
)

如需了解详情,请参阅导出 PyTorch 的模型工件以及有关使用预构建的容器提供 PyTorch 模型的 Jupyter 笔记本。

自定义容器

对于自定义容器,您的模型不必是 TensorFlow 模型,但必须针对 TPU 进行优化。如需了解如何生成针对 TPU 优化的模型,请参阅以下通用机器学习框架的指南:

如需了解如何在 Cloud TPU v5e 上提供使用 JAX、TensorFlow 或 PyTorch 训练的模型,请参阅 Cloud TPU v5e 推理

确保您的自定义容器符合自定义容器要求

您必须提高锁定的内存限制,以便驱动程序可以通过直接内存访问 (DMA) 与 TPU 芯片进行通信。例如:

命令行

ulimit -l 68719476736

Python

import resource

resource.setrlimit(
    resource.RLIMIT_MEMLOCK,
    (
        68_719_476_736_000,  # soft limit
        68_719_476_736_000,  # hard limit
    ),
  )

然后,请参阅使用自定义容器进行预测,了解如何使用自定义容器导入模型。如果您要实现预处理或后处理逻辑,请考虑使用自定义预测例程

创建端点

为 Cloud TPU 创建端点的说明与创建任何端点的说明相同。

例如,以下命令会创建一个 endpoint 资源:

endpoint = aiplatform.Endpoint.create(display_name='My endpoint')

响应中包含新端点的 ID,您将在后续步骤中用到它。

如需详细了解如何创建端点,请参阅将模型部署到端点

部署模型

将模型部署到 Cloud TPU 的说明与部署任何模型的说明相同,除了指定以下受支持的 Cloud TPU 机器类型之一:

机器类型 TPU 芯片数量
ct5lp-hightpu-1t 1
ct5lp-hightpu-4t 4
ct5lp-hightpu-8t 8

TPU 加速器内置于机器类型中。您无需指定加速器类型或加速器数量。

例如,以下命令通过调用 deployModel 部署模型:

machine_type = 'ct5lp-hightpu-1t'

deployed_model = model.deploy(
    endpoint=endpoint,
    deployed_model_display_name='My deployed model',
    machine_type=machine_type,
    traffic_percentage=100,
    min_replica_count=1
    sync=True,
)

如需了解详情,请参阅将模型部署到端点

获取在线预测结果

从 Cloud TPU 获取在线预测结果的说明与获取在线预测结果的说明相同。

例如,以下命令通过调用 predict 发送在线预测请求:

deployed_model.predict(...)

对于自定义容器,请参阅自定义容器的预测请求和响应要求

保障容量

默认情况下,Custom model serving TPU v5e cores per region 的配额为 0。

如需申请增加配额,请参阅申请更高的配额限制

价格

TPU 机器类型按小时计费,就像 Vertex Prediction 中的所有其他机器类型一样。如需了解详情,请参阅预测价格

后续步骤