TensorFlow コードを TPU Pod スライスで実行する
このドキュメントでは、TPU Pod で TensorFlow を使用して簡単な計算を行う方法を説明します。次の手順を行います。
- TensorFlow ソフトウェアを使用して TPU Pod スライスを作成する
- SSH を使用して TPU VM に接続する
- 単純なスクリプトを作成して実行する
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
簡単な計算スクリプトを作成して実行する
次の環境変数を設定します。
(vm)$ export TPU_NAME=tpu-name (vm)$ export TPU_LOAD_LIBRARY=0
現在のディレクトリに
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)
次のコマンドでこのスクリプトを実行します。
(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 の使用を終了したら、次の手順に沿ってリソースをクリーンアップします。
Compute Engine から切断します。
(vm)$ exit
Cloud TPU を削除します。
$ gcloud compute tpus tpu-vm delete tpu-name \ --zone europe-west4-a
次のコマンドを実行して、リソースが削除されたことを確認します。TPU がリストに表示されないことを確認します。削除には数分かかることがあります。
$ gcloud compute tpus tpu-vm list \ --zone europe-west4-a