このページは Cloud Translation API によって翻訳されました。
Switch to English

使用のTPU

TensorFlow.org上に表示します Googleのコラボで実行します GitHubの上のソースを表示 ダウンロードノート

クラウドのTPUのための実験的なサポートがKerasとGoogleのコラボのために現在利用可能です。ランタイム>変更の実行時の型>ハードウェア・アクセラレータ> TPU:あなたはこのコラボノートPCを実行する前に、ハードウェアアクセラレータは、お使いのノートブックの設定をチェックしてTPUであることを確認してください。

 import tensorflow as tf

import os
import tensorflow_datasets as tfds
 

TPUの初期化

TPUは、ユーザのpythonプログラムを実行しているローカルのプロセスは異なっているクラウドTPUの労働者に通常あります。したがって、いくつかの初期化作業は、リモートクラスタに接続したTPUを初期化するために行われる必要があります。注意tpu引数TPUClusterResolverちょうどコラボのために特別なアドレスです。 Google Compute Engineの(GCE)上で実行されている場合には、あなたの代わりにあなたのCloudTPUの名前を渡す必要があります。

 resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_cluster(resolver)
# This is the TPU initialization code that has to be at the beginning.
tf.tpu.experimental.initialize_tpu_system(resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))
 
INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.10:8470

INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.10:8470

INFO:tensorflow:Clearing out eager caches

INFO:tensorflow:Clearing out eager caches

INFO:tensorflow:Finished initializing TPU system.

INFO:tensorflow:Finished initializing TPU system.

All devices:  [LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:7', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:6', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:5', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:4', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:0', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:1', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:2', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:3', device_type='TPU')]

マニュアルデバイスの配置

TPUが初期化された後、あなたは、単一のTPUデバイス上の計算を配置するために、手動デバイスの配置を使用することができます。

 a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
with tf.device('/TPU:0'):
  c = tf.matmul(a, b)
print("c device: ", c.device)
print(c)
 
c device:  /job:worker/replica:0/task:0/device:TPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

流通戦略

ほとんどの時間のユーザーは、データ並列の方法で複数のTPU上でモデルを実行します。配布戦略はCPU、GPUの又はのTPUにモデルを駆動するために使用することができる抽象化です。単に流通戦略をスワップアウトし、モデルは、特定のデバイス上で実行されます。参照流通戦略ガイドを詳細については。

まず、作成TPUStrategyオブジェクトを。

 strategy = tf.distribute.experimental.TPUStrategy(resolver)
 
INFO:tensorflow:Found TPU system:

INFO:tensorflow:Found TPU system:

INFO:tensorflow:*** Num TPU Cores: 8

INFO:tensorflow:*** Num TPU Cores: 8

INFO:tensorflow:*** Num TPU Workers: 1

INFO:tensorflow:*** Num TPU Workers: 1

INFO:tensorflow:*** Num TPU Cores Per Worker: 8

INFO:tensorflow:*** Num TPU Cores Per Worker: 8

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)

それは、すべてのTPUコアで実行できるように、計算を複製するには、単純にそれを渡すことができstrategy.run API。以下の例では、すべてのコアが同じ入力を取得することである(a, b) 、独立して各コアにMATMULを行います。出力は、すべてのレプリカからの値になります。

 @tf.function
def matmul_fn(x, y):
  z = tf.matmul(x, y)
  return z

z = strategy.run(matmul_fn, args=(a, b))
print(z)
 
PerReplica:{
  0: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  1: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  2: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  3: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  4: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  5: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  6: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  7: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)
}

TPUは上の分類

私たちは基本的な概念を学んできたように、より具体的な例を見て時間です。このガイドでは、流通戦略の使用方法を示しtf.distribute.experimental.TPUStrategyクラウドTPUを駆動し、Kerasモデルを訓練します。

Kerasモデルを定義します。

以下は、CPUやGPU上で使用するものと変わらないKerasを使用してMNISTモデルの定義は、あります。 Kerasモデルの作成が内側にする必要があることを注意strategy.scope変数は、各TPUデバイス上に作成することができますので、。コードの他の部分は、戦略スコープ内である必要はありません。

 def create_model():
  return tf.keras.Sequential(
      [tf.keras.layers.Conv2D(256, 3, activation='relu', input_shape=(28, 28, 1)),
       tf.keras.layers.Conv2D(256, 3, activation='relu'),
       tf.keras.layers.Flatten(),
       tf.keras.layers.Dense(256, activation='relu'),
       tf.keras.layers.Dense(128, activation='relu'),
       tf.keras.layers.Dense(10)])
 

入力データセット

効率的な使用tf.data.DatasetクラウドTPUを使用しているとき、あなたが十分にすぐにそれらのデータを送ることができない限り、クラウドのTPUを使用することは不可能であるとして、APIは、非常に重要です。参照してください入力パイプラインパフォーマンスガイドをデータセットのパフォーマンスの詳細については。

すべてが、(使用する最も簡単な実験のためにtf.data.Dataset.from_tensor_slicesまたはその他でグラフデータを)すべてのデータファイルを保存する必要がありますGoogleのクラウドストレージ(GCS)バケットにデータセットで読み取ります。

ほとんどのユースケースでは、あなたにデータを変換することをお勧めしTFRecord形式と使用tf.data.TFRecordDatasetそれを読みます。参照TFRecordとtf.Exampleチュートリアルこれを行う方法の詳細については、を。これは、しかし、ハード要件ではなく、あなたは他のデータセットの読者(使用することができますFixedLengthRecordDatasetTextLineDatasetご希望の場合)。

小さなデータセットを使用してメモリに完全にロードすることができtf.data.Dataset.cache

関係なく使用されるデータフォーマットの、強くあなたが100メガバイトのために、大きなファイルを使用することをお勧めします。ファイルを開くのオーバーヘッドがかなり高いので、これは、このネットワーク接続の設定では特に重要です。

ここでは、使用する必要がありますtensorflow_datasets MNISTトレーニングデータのコピーを取得するためのモジュール。ことを注意try_gcs公共GCSバケットで利用可能なコピーを使用するように指定されています。これを指定しない場合は、TPUは、ダウンロードされたデータにアクセスすることはできません。

 def get_dataset(batch_size, is_training=True):
  split = 'train' if is_training else 'test'
  dataset, info = tfds.load(name='mnist', split=split, with_info=True,
                            as_supervised=True, try_gcs=True)

  def scale(image, label):
    image = tf.cast(image, tf.float32)
    image /= 255.0

    return image, label

  dataset = dataset.map(scale)

  # Only shuffle and repeat the dataset in training. The advantage to have a
  # infinite dataset for training is to avoid the potential last partial batch
  # in each epoch, so users don't need to think about scaling the gradients
  # based on the actual batch size.
  if is_training:
    dataset = dataset.shuffle(10000)
    dataset = dataset.repeat()

  dataset = dataset.batch(batch_size)

  return dataset
 

Keras高レベルAPIを使用してモデルをトレーニング

あなたは、単にKerasフィット/コンパイルAPIとモデルを訓練することができます。ここでは何もTPUの特定ではありませんあなたはmutlipleのGPUを持っていた場所を使用する場合は、下の同じコードを記述しますMirroredStrategyではなくTPUStrategy 。より多くを学ぶために、チェックアウトKerasを持つ分散型トレーニングチュートリアルを。

 with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['sparse_categorical_accuracy'])

batch_size = 200
steps_per_epoch = 60000 // batch_size

train_dataset = get_dataset(batch_size, is_training=True)
test_dataset = get_dataset(batch_size, is_training=False)

model.fit(train_dataset,
          epochs=5,
          steps_per_epoch=steps_per_epoch,
          validation_data=test_dataset)
 
Epoch 1/5
300/300 [==============================] - 10s 33ms/step - sparse_categorical_accuracy: 0.9592 - loss: 0.1339 - val_loss: 0.0485 - val_sparse_categorical_accuracy: 0.9833
Epoch 2/5
300/300 [==============================] - 8s 26ms/step - sparse_categorical_accuracy: 0.9895 - loss: 0.0348 - val_loss: 0.0402 - val_sparse_categorical_accuracy: 0.9884
Epoch 3/5
300/300 [==============================] - 8s 25ms/step - sparse_categorical_accuracy: 0.9938 - loss: 0.0194 - val_loss: 0.0438 - val_sparse_categorical_accuracy: 0.9871
Epoch 4/5
300/300 [==============================] - 7s 25ms/step - sparse_categorical_accuracy: 0.9956 - loss: 0.0129 - val_loss: 0.0455 - val_sparse_categorical_accuracy: 0.9869
Epoch 5/5
300/300 [==============================] - 8s 25ms/step - sparse_categorical_accuracy: 0.9964 - loss: 0.0110 - val_loss: 0.0380 - val_sparse_categorical_accuracy: 0.9889

<tensorflow.python.keras.callbacks.History at 0x7fc5d0329748>

カスタムトレーニングループを使用してモデルを訓練します。

また、使用してモデルを作成し、訓練することができtf.functiontf.distribute直接APIを。 strategy.experimental_distribute_datasets_from_function APIは、データセットの機能与えられたデータセットを配布するために使用されます。データセットに渡されたバッチサイズは、この場合、レプリカバッチサイズの代わりに、グローバルなバッチサイズごとになることに注意してください。より多くを学ぶために、チェックアウトtf.distribute.Strategyを持つカスタムトレーニングチュートリアルを。

まず、モデル、データセットとtf.functionsを作成します。

 # Create the model, optimizer and metrics inside strategy scope, so that the
# variables can be mirrored on each device.
with strategy.scope():
  model = create_model()
  optimizer = tf.keras.optimizers.Adam()
  training_loss = tf.keras.metrics.Mean('training_loss', dtype=tf.float32)
  training_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
      'training_accuracy', dtype=tf.float32)

# Calculate per replica batch size, and distribute the datasets on each TPU
# worker.
per_replica_batch_size = batch_size // strategy.num_replicas_in_sync

train_dataset = strategy.experimental_distribute_datasets_from_function(
    lambda _: get_dataset(per_replica_batch_size, is_training=True))

@tf.function
def train_step(iterator):
  """The step function for one training step"""

  def step_fn(inputs):
    """The computation to run on each TPU device."""
    images, labels = inputs
    with tf.GradientTape() as tape:
      logits = model(images, training=True)
      loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits, from_logits=True)
      loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))
    training_loss.update_state(loss * strategy.num_replicas_in_sync)
    training_accuracy.update_state(labels, logits)

  strategy.run(step_fn, args=(next(iterator),))
 

そして、トレーニングループを実行します。

 steps_per_eval = 10000 // batch_size

train_iterator = iter(train_dataset)
for epoch in range(5):
  print('Epoch: {}/5'.format(epoch))

  for step in range(steps_per_epoch):
    train_step(train_iterator)
  print('Current step: {}, training loss: {}, accuracy: {}%'.format(
      optimizer.iterations.numpy(),
      round(float(training_loss.result()), 4),
      round(float(training_accuracy.result()) * 100, 2)))
  training_loss.reset_states()
  training_accuracy.reset_states()
 
Epoch: 0/5
Current step: 300, training loss: 0.1348, accuracy: 95.77%
Epoch: 1/5
Current step: 600, training loss: 0.0345, accuracy: 98.95%
Epoch: 2/5
Current step: 900, training loss: 0.0187, accuracy: 99.44%
Epoch: 3/5
Current step: 1200, training loss: 0.0131, accuracy: 99.57%
Epoch: 4/5
Current step: 1500, training loss: 0.0065, accuracy: 99.8%

内の複数の工程により、パフォーマンスの向上tf.function

性能は内で複数のステップを実行することによって改善することができるtf.function 。これは、ラップすることによって達成されるstrategy.runでコールをtf.range内部tf.function 、オートグラフはに変換されますtf.while_loop TPUワーカーに。

より良い性能を持つものの、内部の単一のステップと比較トレードオフがあるtf.function 。中に複数のステップ実行tf.functionあまり柔軟性がある、あなたは熱心に物事を実行するか、またはステップ内の任意のPythonコードすることはできません。

 @tf.function
def train_multiple_steps(iterator, steps):
  """The step function for one training step"""

  def step_fn(inputs):
    """The computation to run on each TPU device."""
    images, labels = inputs
    with tf.GradientTape() as tape:
      logits = model(images, training=True)
      loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits, from_logits=True)
      loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))
    training_loss.update_state(loss * strategy.num_replicas_in_sync)
    training_accuracy.update_state(labels, logits)

  for _ in tf.range(steps):
    strategy.run(step_fn, args=(next(iterator),))

# Convert `steps_per_epoch` to `tf.Tensor` so the `tf.function` won't get 
# retraced if the value changes.
train_multiple_steps(train_iterator, tf.convert_to_tensor(steps_per_epoch))

print('Current step: {}, training loss: {}, accuracy: {}%'.format(
      optimizer.iterations.numpy(),
      round(float(training_loss.result()), 4),
      round(float(training_accuracy.result()) * 100, 2)))
 
Current step: 1800, training loss: 0.0084, accuracy: 99.74%

次のステップ