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.18.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 にアクセスできることを確認する
現在のディレクトリに
tpu-count.py
という名前のファイルを作成し、次のスクリプトをコピーして貼り付けます。import tensorflow as tf print(f"TensorFlow can access {len(tf.config.list_logical_devices('TPU'))} TPU cores")
スクリプトを実行します。
(vm)$ python3 tpu-count.py
スクリプトの出力には、TPU VM で使用可能な TPU コアの数が表示されます。
TensorFlow can access 8 TPU cores
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 アカウントに課金されないようにするには、次の手順を行います。
Compute Engine インスタンスとの接続を切断していない場合は切断します。
(vm)$ exit
プロンプトが
username@projectname
に変わります。これは、現在、Cloud Shell 内にいることを示しています。Cloud TPU を削除します。
$ gcloud compute tpus tpu-vm delete tpu-name \ --zone=europe-west4-a
gcloud compute tpus tpu-vm list
を実行して、リソースが削除されたことを確認します。削除には数分かかることがあります。$ gcloud compute tpus tpu-vm list --zone=europe-west4-a
次のステップ
Cloud TPU の詳細については、以下をご覧ください。