Google Cloud 提供对称为张量处理单元 (TPU) 的定制机器学习加速器的访问权限。TPU 经过优化,可以加快机器学习模型的训练和推理速度,非常适合包括自然语言处理、计算机视觉和语音识别在内的各种应用。
本页面介绍了如何将模型部署到单主机 Cloud TPU v5e 中,以便在 Vertex AI 中进行在线预测。
仅支持 Cloud TPU v5e。不支持其他 Cloud TPU 世代。
导入模型
如需在 Cloud TPU 上进行部署,您必须将模型导入 Vertex AI,并将其配置为使用以下容器之一:
- 预构建的优化 TensorFlow 运行时容器
nightly
版本,或者2.15
版或更高版本 - 预构建的 PyTorch TPU 容器
2.1
版或更高版本 - 您自己的支持 TPU 的自定义容器
预构建的优化 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 中的所有其他机器类型一样。如需了解详情,请参阅预测价格。
后续步骤
- 了解如何获取在线预测结果。