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

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

始める前に

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

gcloud を使用した Cloud TPU VM の作成

gcloud コマンドを使用して Cloud TPU を作成します。

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=europe-west4-a
    --accelerator-type=v3-8
    --version=tpu-vm-tf-2.17.0-pjrt
    --project=your-gcp-project-name

コマンドフラグの説明

tpu-name
作成する Cloud TPU の名前。
zone
Cloud TPU を作成するゾーン。
accelerator-type
アクセラレータ タイプでは、作成する Cloud TPU のバージョンとサイズを指定します。 TPU の各バージョンでサポートされているアクセラレータ タイプの詳細については、TPU のバージョンをご覧ください。
version
TPU ランタイムのバージョン。このクイックスタートで使用するバージョンには、TensorFlow がプリインストールされています。
project
Cloud TPU を作成する Google Cloud CLI プロジェクトの名前。

gcloud コマンドの詳細については、gcloud リファレンスをご覧ください。

Cloud TPU VM に接続する

SSH を使用して TPU VM に接続します。

  $ gcloud compute tpus tpu-vm ssh tpu-name 

      --zone europe-west4-a 

      --project=your-gcp-project-name

TensorFlow を使用してサンプルを実行する

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)

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

(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 を削除します。

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

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

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

次のステップ

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