將模型部署至 Cloud TPU VM

Google Cloud 可存取稱為 Tensor Processing Unit (TPU) 的自訂設計機器學習加速器。TPU 經過最佳化，可加速機器學習模型的訓練和推論，因此非常適合用於各種應用程式，包括自然語言處理、電腦視覺和語音辨識。

本頁說明如何將模型部署到 單一主機 Cloud TPU v5e，以便在 Vertex AI 中進行線上推論。

僅支援 Cloud TPU v5e 版。系統不支援其他 Cloud TPU 代別。

如要瞭解 Cloud TPU 第 5 代第 5 版在哪些地點提供，請參閱地點。

匯入模型

如要在 Cloud TPU 上部署模型，您必須將模型匯入 Vertex AI，並設定使用下列其中一個容器：

預先建構的最佳化 TensorFlow 執行階段容器

如要在 Cloud TPU 上匯入及執行 TensorFlow 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 模型的教學課程筆記本。

自訂容器

如果是自訂容器，模型不一定要是 TensorFlow 模型，但必須經過 TPU 最佳化。如要瞭解如何產生 TPU 最佳化模型，請參閱常見 ML 架構的下列指南：

如要瞭解如何在 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。

如要申請提高配額，請參閱申請調整配額。

定價

與 Vertex Prediction 中的所有其他機器類型一樣，TPU 機器類型也是以每小時為單位計費。詳情請參閱「預測定價」。

後續步驟