TensorFlow コードを TPU Pod スライスで実行する

このドキュメントでは、TPU Pod で TensorFlow を使用して簡単な計算を行う方法を説明します。次の手順を行います。

  1. TensorFlow ソフトウェアを使用して TPU Pod スライスを作成する
  2. SSH を使用して TPU VM に接続する
  3. 単純なスクリプトを作成して実行する

TPU VM は、Cloud TPU API を呼び出す権限をサービス アカウントに依存します。デフォルトでは、TPU VM はデフォルトの Compute Engine サービス アカウントを使用します。これには、必要なすべての Cloud TPU 権限が含まれています。独自のサービス アカウントを使用する場合は、TPU 閲覧者のロールをサービス アカウントに追加する必要があります。Google Cloud のロールの詳細については、ロールについてをご覧ください。TPU VM を作成するときに、--service-account フラグを使用して独自のサービス アカウントを指定できます。

TensorFlow ランタイムを使用して v3-32 TPU Pod スライスを作成する

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=europe-west4-a \
  --accelerator-type=v3-32 \
  --version=tpu-vm-tf-2.17.0-pod-pjrt

コマンドフラグの説明

zone
Cloud TPU を作成するゾーン。
accelerator-type
アクセラレータ タイプでは、作成する Cloud TPU のバージョンとサイズを指定します。TPU のバージョンごとにサポートされているアクセラレータ タイプの詳細については、TPU のバージョンをご覧ください。
version
Cloud TPU ソフトウェアのバージョン。

SSH を使用して Cloud TPU VM に接続する

$ gcloud compute tpus tpu-vm ssh tpu-name \
      --zone europe-west4-a

簡単な計算スクリプトを作成して実行する

  1. 次の環境変数を設定します。

    (vm)$ export TPU_NAME=tpu-name
(vm)$ export TPU_LOAD_LIBRARY=0

  2. 現在のディレクトリに tpu-test.py という名前のファイルを作成し、次のスクリプトをコピーして貼り付けます。

    import tensorflow as tf
print("Tensorflow version " + tf.__version__)

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
print('Running on TPU ', cluster_resolver.cluster_spec().as_dict()['worker'])

tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
strategy = tf.distribute.experimental.TPUStrategy(cluster_resolver)

@tf.function
def add_fn(x,y):
z = x + y
return z

x = tf.constant(1.)
y = tf.constant(1.)
z = strategy.run(add_fn, args=(x,y))
print(z)

  3. 次のコマンドでこのスクリプトを実行します。

    (vm)$ python3 tpu-test.py

    このスクリプトは、TPU の各 TensorCore で簡単な計算を実行します。出力は次のようになります。

    PerReplica:{
0: tf.Tensor(2.0, shape=(), dtype=float32),
1: tf.Tensor(2.0, shape=(), dtype=float32),
2: tf.Tensor(2.0, shape=(), dtype=float32),
3: tf.Tensor(2.0, shape=(), dtype=float32),
4: tf.Tensor(2.0, shape=(), dtype=float32),
5: tf.Tensor(2.0, shape=(), dtype=float32),
6: tf.Tensor(2.0, shape=(), dtype=float32),
7: tf.Tensor(2.0, shape=(), dtype=float32),
8: tf.Tensor(2.0, shape=(), dtype=float32),
9: tf.Tensor(2.0, shape=(), dtype=float32),
10: tf.Tensor(2.0, shape=(), dtype=float32),
11: tf.Tensor(2.0, shape=(), dtype=float32),
12: tf.Tensor(2.0, shape=(), dtype=float32),
13: tf.Tensor(2.0, shape=(), dtype=float32),
14: tf.Tensor(2.0, shape=(), dtype=float32),
15: tf.Tensor(2.0, shape=(), dtype=float32),
16: tf.Tensor(2.0, shape=(), dtype=float32),
17: tf.Tensor(2.0, shape=(), dtype=float32),
18: tf.Tensor(2.0, shape=(), dtype=float32),
19: tf.Tensor(2.0, shape=(), dtype=float32),
20: tf.Tensor(2.0, shape=(), dtype=float32),
21: tf.Tensor(2.0, shape=(), dtype=float32),
22: tf.Tensor(2.0, shape=(), dtype=float32),
23: tf.Tensor(2.0, shape=(), dtype=float32),
24: tf.Tensor(2.0, shape=(), dtype=float32),
25: tf.Tensor(2.0, shape=(), dtype=float32),
26: tf.Tensor(2.0, shape=(), dtype=float32),
27: tf.Tensor(2.0, shape=(), dtype=float32),
28: tf.Tensor(2.0, shape=(), dtype=float32),
29: tf.Tensor(2.0, shape=(), dtype=float32),
30: tf.Tensor(2.0, shape=(), dtype=float32),
31: tf.Tensor(2.0, shape=(), dtype=float32)
}

クリーンアップ

TPU VM の使用を終了したら、次の手順に沿ってリソースをクリーンアップします。

  1. Compute Engine から切断します。

    (vm)$ exit

  2. Cloud TPU を削除します。

    $ gcloud compute tpus tpu-vm delete tpu-name \
  --zone europe-west4-a

  3. 次のコマンドを実行して、リソースが削除されたことを確認します。TPU がリストに表示されないことを確認します。削除には数分かかることがあります。

    $ gcloud compute tpus tpu-vm list \
  --zone europe-west4-a