Cloud TPU 推論の概要

サービス提供とは、トレーニング済み機械学習モデルを本番環境にデプロイするプロセスであり、本番環境を使用して推論に使用できます。推論は TPU v5e 以降のバージョンでサポートされています。レイテンシ SLO は提供の優先事項です。

このドキュメントでは、単一ホストの TPU でモデルを提供する方法について説明します。チップ数が 8 個以下の TPU スライスには 1 つの TPU VM またはホストがあり、単一ホストの TPU と呼ばれます。

始める

Cloud TPU を使用するには、 Google Cloud アカウントとプロジェクトが必要です。詳細については、Cloud TPU 環境を設定するをご覧ください。

TPU でサービングするには、次の割り当てをリクエストする必要があります。

• オンデマンド v5e リソース: TPUv5 lite pod cores for serving per project per zone
• プリエンプティブル v5e リソース: Preemptible TPU v5 lite pod cores for serving per project per zone
• オンデマンドの v6e リソース: TPUv6 cores per project per zone
• プリエンプティブル v6e リソース: Preemptible TPUv6 cores per project per zone

TPU 割り当ての詳細については、TPU 割り当てをご覧ください。

JAX モデルの推論とサービス提供

TPU VM でモデルを提供する手順は次のとおりです。

1. TensorFlow の SavedModel 形式でモデルをシリアル化する
2. Inference Converter を使用して、保存したモデルをサービス提供用に準備する
3. TensorFlow Serving を使用してモデルを提供する

SavedModel 形式

SavedModel には、（トレーニング済みパラメータと計算を含む）完全な TensorFlow プログラムが含まれます。元のモデル構築コードを実行する必要はありません。

モデルが JAX で記述されている場合は、jax2tf を使用してモデルを SavedModel 形式でシリアル化する必要があります。

Inference Converter

Cloud TPU Inference Converter は、SavedModel 形式でエクスポートされたモデルを TPU 推論用に準備し、最適化します。Inference Converter は、ローカルシェルまたは TPU VM で実行できます。コンバータの実行に必要なすべてのコマンドライン ツールが含まれているため、TPU VM シェルの使用をおすすめします。Inference Converter の詳細については、Inference Converter ユーザーガイドをご覧ください。

Inference Converter の要件

1. モデルは、TensorFlow または JAX から SavedModel 形式でエクスポートする必要があります。

2. TPU 関数の関数エイリアスを定義する必要があります。詳細については、Inference Converter ユーザーガイドをご覧ください。このガイドの例では、TPU 関数のエイリアスとして tpu_func を使用します。

3. Tensorflow ライブラリ（Cloud TPU Inference Converter の依存関係）は AVX 命令を使用するようにコンパイルされるため、マシンの CPU が Advanced Vector eXtensions（AVX）命令をサポートしていることを確認してください。ほとんどの CPU が AVX をサポートします。

JAX モデルの推論とサービス提供

このセクションでは、jax2tf と TensorFlow Serving を使用して JAX モデルを提供する方法について説明します。

1. jax2tf を使用してモデルを SavedModel 形式にシリアル化する
2. Inference Converter を使用して、サービス提供用の保存済みモデルを準備する
3. TensorFlow Serving を使用してモデルを提供する

jax2tf を使用して JAX モデルを SavedModel 形式にシリアル化する

次の Python 関数は、モデルコード内で jax2tf を使用する方法を示しています。

# Inference function
def model_jax(params, inputs):
  return params[0] + params[1] * inputs

# Wrap the parameter constants as tf.Variables; this will signal to the model
# saving code to save those constants as variables, separate from the
# computation graph.
params_vars = tf.nest.map_structure(tf.Variable, params)

# Build the prediction function by closing over the `params_vars`. If you
# instead were to close over `params` your SavedModel would have no variables
# and the parameters will be included in the function graph.
prediction_tf = lambda inputs: jax2tf.convert(model_jax)(params_vars, inputs)

my_model = tf.Module()
# Tell the model saver what the variables are.
my_model._variables = tf.nest.flatten(params_vars)
my_model.f = tf.function(prediction_tf, jit_compile=True, autograph=False)
tf.saved_model.save(my_model)

jax2tf の詳細については、JAX と Cloud TPU の相互運用をご覧ください。

Inference Converter を使用して、保存したモデルをサービス提供用に準備する

Inference Converter の使用方法については、Inference Converter ガイドをご覧ください。

TensorFlow Serving を使用する

TensorFlow Serving の使用方法については、TensorFlow Serving をご覧ください。

JAX モデル サービング例の前提条件

1. 現在のユーザーを docker グループに追加します。

   sudo usermod -a -G docker ${USER}
   

2. docker をアクティブ グループとして新しいシェル セッションを作成します。

   newgrp docker
   

3. Artifact Registry で認証するように Docker を構成します。

   gcloud auth configure-docker us-docker.pkg.dev
   

4. Inference Converter Docker イメージをダウンロードします。

   docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
   

5. Cloud TPU Serving Docker イメージをダウンロードします。

   docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
   

6. SSH を使用して TPU VM に接続し、推論デモコードをインストールします。

   gcloud storage cp \
   "gs://cloud-tpu-inference-public/demo" \
   . \
   --recursive
   

7. JAX デモの依存関係をインストールします。

   pip install -r ./demo/jax/requirements.txt
   

推論用に JAX BERT モデルを提供する

事前トレーニング済みの BERT モデルは Hugging Face からダウンロードできます。

1. Flax BERT モデルから TPU 互換の TensorFlow 保存モデルをエクスポートします。

   cd demo/jax/bert
   python3 export_bert_model.py
   

2. Cloud TPU モデルサーバー コンテナを起動します。

   docker run -t --rm --privileged -d \
     -p 8500:8500 -p 8501:8501 \
     --mount type=bind,source=/tmp/jax/bert_tpu,target=/models/bert \
     -e MODEL_NAME=bert \
     us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
   

   コンテナの起動から約 30 秒後に、モデルサーバー コンテナのログを確認し、gRPC サーバーと HTTP サーバーが稼働していることを確認します。

   CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
   docker logs ${CONTAINER_ID}
   

   次の情報で終わるログエントリが表示されている場合、サーバーはリクエストを処理する準備ができています。

   2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ...
   [warn] getaddrinfo: address family for nodename not supported
   2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ...
   [evhttp_server.cc : 245] NET_LOG: Entering the event loop ...
   

3. モデルサーバーに推論リクエストを送信します。

   python3 bert_request.py
   

   出力は次のようになります。

   For input "The capital of France is [MASK].", the result is ". the capital of france is paris.."
   For input "Hello my name [MASK] Jhon, how can I [MASK] you?", the result is ". hello my name is jhon, how can i help you?."
   

4. クリーンアップする。

   他のデモを実行する前に、Docker コンテナをクリーンアップしてください。

   CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
   docker stop ${CONTAINER_ID}
   

   モデル アーティファクトをクリーンアップします。

   sudo rm -rf /tmp/jax/
   

推論用に JAX Stable Diffusion を提供する

事前トレーニング済みの Stable Diffusion モデルは、Hugging Face からダウンロードできます。

1. TPU 互換の TF2 保存モデル形式で Stable Diffusion モデルをダウンロードします。

   cd demo/jax/stable_diffusion
   python3 export_stable_diffusion_model.py
   

2. モデル用の Cloud TPU モデルサーバー コンテナを起動します。

   docker run -t --rm --privileged -d \
     -p 8500:8500 -p 8501:8501 \
     --mount type=bind,source=/tmp/jax/stable_diffusion_tpu,target=/models/stable_diffusion \
     -e MODEL_NAME=stable_diffusion \
     us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
   

   2 分ほど経過したら、モデルサーバー コンテナのログを確認して、gRPC サーバーと HTTP サーバーが実行されていることを確認します。

   CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
   docker logs ${CONTAINER_ID}
   

   次の情報で終わるログが表示されている場合、サーバーはリクエストを処理する準備ができています。

   2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ...
   [warn] getaddrinfo: address family for nodename not supported
   2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ...
   [evhttp_server.cc : 245] NET_LOG: Entering the event loop ...
   

3. モデルサーバーにリクエストを送信します。

   python3 stable_diffusion_request.py
   

   このスクリプトは、「Painting of a squirrel skating in New York」をプロンプトとして送信します。出力画像は現在のディレクトリに stable_diffusion_images.jpg として保存されます。

4. クリーンアップする。

   他のデモを実行する前に、Docker コンテナをクリーンアップしてください。

   CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
   docker stop ${CONTAINER_ID}
   

   モデル アーティファクトをクリーンアップする

   sudo rm -rf /tmp/jax/
   

TensorFlow Serving

次の手順では、TPU VM で TensorFlow モデルを提供する方法を示します。

TensorFlow Serving ワークフロー

1. TPU VM 用の TensorFlow Serving Docker イメージをダウンロードします。

   サンプル環境変数を設定します。

   export YOUR_LOCAL_MODEL_PATH=model-path
   export MODEL_NAME=model-name
   export IMAGE_NAME=us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

   Docker イメージをダウンロードします。

   docker pull ${IMAGE_NAME}
   

2. 現在のユーザーを docker グループに追加します。

   sudo usermod -a -G docker ${USER}
   

3. docker をアクティブ グループとして新しいシェル セッションを作成します。

   newgrp docker
   

4. Artifact Registry で認証するように Docker を構成します。

   gcloud auth configure-docker us-docker.pkg.dev
   

5. Inference Converter Docker イメージをダウンロードします。

   docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
   

6. Cloud TPU Serving Docker イメージをダウンロードします。

   docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
   

7. デモコードをダウンロードします。

   gcloud storage cp \
   "gs://cloud-tpu-inference-public/demo" \
   . \
   --recursive
   

8. TensorFlow デモの依存関係をインストールします。

   pip install -r ./demo/tf/requirements.txt
   

9. TPU VM で TensorFlow Serving Docker イメージを使用して TensorFlow モデルを提供します。

   # PORT 8500 is for gRPC model server and 8501 is for HTTP/REST model server.
   docker run -t --rm --privileged -d \
     -p 8500:8500 -p 8501:8501 \
     --mount type=bind,source=${YOUR_LOCAL_MODEL_PATH},target=/models/${MODEL_NAME} \
     -e MODEL_NAME=${MODEL_NAME} \
     ${IMAGE_NAME}
   

10. Serving Client API を使用してモデルをクエリします。

    • REST クライアント API
    • gRPC クライアント API

TensorFlow ResNet-50 サービス提供デモを実行する

1. Keras ResNet-50 モデルから TPU 互換の TF2 SavedModel をエクスポートします。

   cd demo/tf/resnet-50
   python3 export_resnet_model.py
   

2. モデル用の TensorFlow モデルサーバー コンテナを起動します。

   docker run -t --rm --privileged -d \
     -p 8500:8500 -p 8501:8501 \
     --mount type=bind,source=/tmp/tf/resnet_tpu,target=/models/resnet \
     -e MODEL_NAME=resnet \
     us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
   

   モデルサーバー コンテナのログを確認し、gRPC サーバーと HTTP サーバーが稼働していることを確認します。

   CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
   docker logs ${CONTAINER_ID}
   

   次の情報で終わるログが表示されている場合、サーバーはリクエストを処理する準備ができています。30 秒ほどかかります。

   2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ...
   [warn] getaddrinfo: address family for nodename not supported
   2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ...
   [evhttp_server.cc : 245] NET_LOG: Entering the event loop ...
   

3. リクエストをモデルサーバーに送信します。

   リクエストされた画像は、https://i.imgur.com/j9xCCzn.jpeg のバナナです。

   python3 resnet_request.py
   

   出力は次のようになります。

   Predict result: [[('n07753592', 'banana', 0.94921875), ('n03532672', 'hook', 0.022338867), ('n07749582', 'lemon', 0.005126953)]]
   

4. クリーンアップする。

   他のデモを実行する前に、Docker コンテナをクリーンアップしてください。

   CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
   docker stop ${CONTAINER_ID}
   

   モデル アーティファクトをクリーンアップします。

   sudo rm -rf /tmp/tf/
   

PyTorch モデルの推論とサービス提供

PyTorch で記述されたモデルの場合、ワークフローは次のとおりです。

1. TorchDynamo と PyTorch/XLA を使用して読み込みと推論を行う Python モデル ハンドラを作成する
2. TorchModelArchiver を使用してモデル アーカイブを作成する
3. TorchServe を使用してモデルをサービングする

TorchDynamo と PyTorch/XLA

TorchDynamo（Dynamo）は、PyTorch プログラムを高速化するように設計された Python レベルの JIT コンパイラです。コンパイラ バックエンドがフックするためのクリーンな API を提供します。実行直前に Python バイトコードを動的に変更します。PyTorch/XLA 2.0 リリースには、Dynamo を使用して推論とトレーニングを行うための試験運用版のバックエンドがあります。

Dynamo は、モデルパターンを認識するときに Torch FX（FX）グラフを提供し、PyTorch/XLA は lazy tensor アプローチを使用して FX グラフをコンパイルし、コンパイル済みの関数を返します。Dynamo の詳細については、以下をご覧ください。

• Pytorch Dev Discussions の投稿
• TorchDynamo のドキュメント
• 詳しくは、PyTorch 2.0 と XLAをご覧ください

torch.compile を使用して densenet161 推論を実行する簡単なコードの例を次に示します。

import torch
import torchvision
import torch_xla.core.xla_model as xm

def eval_model(loader):
  device = xm.xla_device()
  xla_densenet161 = torchvision.models.densenet161().to(device)
  xla_densenet161.eval()
  dynamo_densenet161 = torch.compile(
      xla_densenet161, backend='torchxla_trace_once')
  for data, _ in loader:
    output = dynamo_densenet161(data)

TorchServe

提供されている torchserve-tpu Docker イメージを使用して、アーカイブされた PyTorch モデルを TPU VM で提供できます。

1. 現在のユーザーを docker グループに追加します。

   sudo usermod -a -G docker ${USER}
   

2. docker をアクティブ グループとして新しいシェル セッションを作成します。

   newgrp docker
   

3. Artifact Registry で認証するように Docker を構成します。

   gcloud auth configure-docker us-docker.pkg.dev
   

4. Cloud TPU TorchServe Docker イメージを TPU VM にダウンロードします。

   CLOUD_TPU_TORCHSERVE_IMAGE_URL=us-docker.pkg.dev/cloud-tpu-images/inference/torchserve-tpu:v0.9.0-2.1.0
   docker pull ${CLOUD_TPU_TORCHSERVE_IMAGE_URL}
   

モデル アーティファクトを収集する

始めに、モデルハンドラを指定する必要があります。このハンドラは、TorchServe モデルサーバー ワーカーにモデルの読み込み、入力データの処理、推論の実行を指示します。また、TorchServe デフォルトの推論ハンドラ（ソース）、または base_handler.py に従って独自のカスタムモデル ハンドラを開発します。 トレーニング済みモデルとモデル定義ファイルを指定することも必要になる場合があります。

次の Densenet 161 の例では、モデル アーティファクトと TorchServe が提供するデフォルトの画像分類ハンドラを使用します。

• model.py
• index_to_name.json
• image_classifier.py

1. いくつかの環境変数を構成します。

   CWD="$(pwd)"
   
   WORKDIR="${CWD}/densenet_161"
   
   mkdir -p ${WORKDIR}/model-store
   mkdir -p ${WORKDIR}/logs
   

2. TorchServe 画像分類モデルのサンプルからモデル アーティファクトをダウンロードしてコピーします。

   git clone https://github.com/pytorch/serve.git
   
   cp ${CWD}/serve/examples/image_classifier/densenet_161/model.py ${WORKDIR}
   cp ${CWD}/serve/examples/image_classifier/index_to_name.json ${WORKDIR}
   

3. モデルの重みをダウンロードします。

   wget https://download.pytorch.org/models/densenet161-8d451a50.pth -O densenet161-8d451a50.pth
   
   mv densenet161-8d451a50.pth ${WORKDIR}
   

4. Dynamo バックエンドを使用する TorchServe モデル構成ファイルを作成します。

   echo 'pt2: "torchxla_trace_once"' >> ${WORKDIR}/model_config.yaml
   

   次のファイルとディレクトリが表示されます。

   >> ls ${WORKDIR}
   model_config.yaml
   index_to_name.json
   logs
   model.py
   densenet161-8d451a50.pth
   model-store
   

モデル アーカイブ ファイルを生成する

Cloud TPU TorchServe で PyTorch モデルを提供する場合は、Torch Model Archiver を使用して、モデル ハンドラとすべてのモデル アーティファクトをモデル アーカイブ ファイル (*.mar) にパッケージ化する必要があります。

torch-model-archiver を使用してモデル アーカイブ ファイルを生成します。

MODEL_NAME=Densenet161

docker run \
    --privileged  \
    --shm-size 16G \
    --name torch-model-archiver \
    -it \
    -d \
    --rm \
    --mount type=bind,source=${WORKDIR},target=/home/model-server/ \
    ${CLOUD_TPU_TORCHSERVE_IMAGE_URL} \
    torch-model-archiver \
        --model-name ${MODEL_NAME} \
        --version 1.0 \
        --model-file model.py \
        --serialized-file densenet161-8d451a50.pth \
        --handler image_classifier \
        --export-path model-store \
        --extra-files index_to_name.json \
        --config-file model_config.yaml

model-store ディレクトリに、生成されたモデル アーカイブ ファイルが表示されます。

>> ls ${WORKDIR}/model-store
Densenet161.mar

推論リクエストを処理する

モデル アーカイブ ファイルが作成されたので、TorchServe モデルサーバーを起動して推論リクエストを処理できます。

1. TorchServe モデルサーバーを起動します。

   docker run \
       --privileged  \
       --shm-size 16G \
       --name torchserve-tpu \
       -it \
       -d \
       --rm \
       -p 7070:7070 \
       -p 7071:7071 \
       -p 8080:8080 \
       -p 8081:8081 \
       -p 8082:8082 \
       -p 9001:9001 \
       -p 9012:9012 \
       --mount type=bind,source=${WORKDIR}/model-store,target=/home/model-server/model-store \
       --mount type=bind,source=${WORKDIR}/logs,target=/home/model-server/logs \
       ${CLOUD_TPU_TORCHSERVE_IMAGE_URL} \
       torchserve \
           --start \
           --ncs \
           --models ${MODEL_NAME}.mar \
           --ts-config /home/model-server/config.properties
   

2. モデルサーバーの健全性をクエリします。

   curl http://localhost:8080/ping
   

   モデルサーバーが稼働している場合は、次のように表示されます。

   {
     "status": "Healthy"
   }
   

   現在登録されているモデルのデフォルト バージョンをクエリするには、次のコマンドを使用します。

   curl http://localhost:8081/models
   

   登録済みのモデルが表示されます。

   {
     "models": [
       {
         "modelName": "Densenet161",
         "modelUrl": "Densenet161.mar"
       }
     ]
   }
   

   推論用に画像をダウンロードするには、次のコマンドを使用します。

   curl -O https://raw.githubusercontent.com/pytorch/serve/master/docs/images/kitten_small.jpg
   
   mv kitten_small.jpg ${WORKDIR}
   

   モデルサーバーに推論リクエストを送信するには、次のコマンドを使用します。

   curl http://localhost:8080/predictions/${MODEL_NAME} -T ${WORKDIR}/kitten_small.jpg
   

   次のようなレスポンスが表示されます。

   {
     "tabby": 0.47878125309944153,
     "lynx": 0.20393909513950348,
     "tiger_cat": 0.16572578251361847,
     "tiger": 0.061157409101724625,
     "Egyptian_cat": 0.04997897148132324
   }
   

3. モデルサーバーのログ

   ログにアクセスするには、次のコマンドを使用します。

   ls ${WORKDIR}/logs/
   cat ${WORKDIR}/logs/model_log.log
   

   ログに次のメッセージが表示されます。

   "Compiled model with backend torchxla\_trace\_once"
   

クリーンアップ

Docker コンテナを停止します。

rm -rf serve
rm -rf ${WORKDIR}

docker stop torch-model-archiver
docker stop torchserve-tpu

プロファイリング

推論を設定したら、プロファイラを使用してパフォーマンスと TPU の使用率を分析できます。プロファイリングの詳細については、以下をご覧ください。

• Cloud TPU でのプロファイリング
• TensorFlow プロファイリング
• PyTorch プロファイリング
• JAX プロファイリング