Cloud TPU TensorFlow クイックスタート

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

始める前に

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

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

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

TPU VM

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

コマンドフラグの説明

zone
Cloud TPU を作成するゾーン
accelerator-type
作成する Cloud 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.6.0 \
--accelerator-type=v3-8

コマンドフラグの説明

project
GCP プロジェクト 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 のタイプ

Cloud TPU VM に接続する

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

TPU VM

$ gcloud alpha 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-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(tpu='local')
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 の各コアで簡単な計算を実行します。出力には、各 TPU コアからの出力結果が表示されます。

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 alpha 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 alpha compute tpus tpu-vm list を実行して、リソースが削除されたことを確認します。削除には数分かかることがあります。

    TPU VM

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

    TPU ノード

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

次のステップ

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