将模型部署到 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 中的所有其他机器类型一样。如需了解详情，请参阅预测价格。

后续步骤