TensorFlow を使用して Cloud TPU VM で計算を実行する

このクイックスタートでは、Cloud TPU の作成、TensorFlow のインストール、Cloud TPU で簡単な計算を行う方法を説明します。Cloud TPU でモデルをトレーニングする方法の詳細については、Cloud TPU のチュートリアルをご覧ください。

始める前に

このクイックスタートを使用する前に、Google Cloud Platform アカウントを作成し、Google Cloud CLI をインストールして gcloud コマンドを構成する必要があります。詳細しくは、アカウントと Cloud TPU プロジェクトを設定するをご覧ください。

gcloud を使用して Cloud TPU VM またはノードを作成する

gcloud コマンドを使用して、Compute Engine Cloud TPU を起動します。使用するコマンドは、TPU VM を使用するか TPU ノードを使用するかによって異なります。2 つの VM アーキテクチャの詳細については、システム アーキテクチャをご覧ください。gcloud コマンドの詳細については、gcloud リファレンスをご覧ください。

TPU VM

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

コマンドフラグの説明

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

TPU ノード

$ gcloud compute tpus execution-groups create \
--name=tpu-name \
--zone=europe-west4-a \
--disk-size=300 \
--machine-type=n1-standard-16 \
--tf-version=2.12.0 \

コマンドフラグの説明

project
: Google Cloud プロジェクト ID
name
作成する Cloud TPU の名前。
zone
Cloud TPU を作成するゾーン。
disk-size
gcloud コマンドで作成された VM のハードディスクのサイズ（GB）。
machine-type
作成する Compute Engine VM のマシンタイプ。
tf-version
Tensorflow gcloud のバージョンが VM にインストールされます。
をご覧ください。
accelerator-type
アクセラレータ タイプでは、作成する Cloud TPU のバージョンとサイズを指定します。TPU のバージョンごとにサポートされているアクセラレータ タイプの詳細については、TPU のバージョンをご覧ください。

Cloud TPU VM に接続する

TPU VM を使用する場合は、SSH を使用して明示的に TPU VM に接続する必要があります。 TPU ノードを使用する場合は、Compute Engine VM に自動的に SSH 接続する必要があります。自動的に接続されない場合は、次のコマンドを使用します。

TPU VM

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

TPU ノード

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

TensorFlow を使用した簡単なサンプルの実行

TPU VM

TPU VM に接続したら、次の環境変数を設定します。

  (vm)$ export TPU_NAME=local

TPU を作成するときに、--version パラメータを -pjrt で終わるバージョンに設定した場合は、次の環境変数を設定して PJRT ランタイムを有効にします。

  (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
  (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

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

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

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

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

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

TPU ノード

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

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

tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='your-tpu-name')  # TPU detection
print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])

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

@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)
}

クリーンアップ

このページで使用したリソースについて、Google Cloud アカウントに課金されないようにするには、次の手順を行います。

  1. Compute Engine インスタンスとの接続を切断していない場合は切断します。

    (vm)$ exit

    プロンプトが username@projectname に変わります。これは、現在、Cloud Shell 内にいることを示しています。

  2. Cloud TPU を削除します。

    TPU VM

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

    TPU ノード

    $ gcloud compute tpus execution-groups delete tpu-name \
--zone=europe-west4-a

  3. gcloud compute tpus tpu-vm list を実行して、リソースが削除されたことを確認します。削除には数分かかることがあります。

    TPU VM

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

    TPU ノード

    $ gcloud compute tpus execution-groups list --zone=europe-west4-a

次のステップ

Cloud TPU の詳細については、以下をご覧ください。