コレクションでコンテンツを整理 必要に応じて、コンテンツの保存と分類を行います。

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

TensorFlow を実行する TPU VM Pod を設定して計算を実行する

  1. TensorFlow を実行する Pod を設定して計算を実行するには、次の手順を使用します。

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

    コマンドフラグの説明

    zone
    Cloud TPU を作成するゾーン
    accelerator-type
    作成する Cloud TPU のタイプ
    version
    Cloud TPU ソフトウェアのバージョン。デフォルトは、最新の TensorFlow ソフトウェア バージョンです。
  2. Cloud TPU VM に接続する

    Cloud TPU Pod の TPU VM に SSH 接続します。

    $ gcloud compute tpus tpu-vm ssh tpu-name \
      --zone europe-west4-a
    
  3. 次の環境変数を設定します。

    (vm)$ export TPU_NAME=tpu-name
    (vm)$ export TPU_LOAD_LIBRARY=0
    
  4. 現在のディレクトリに 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 の使用を終了したら、次の手順に沿ってリソースをクリーンアップします。

  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