Google I / Oが5月18〜20日に戻ってきます。スペースを予約してスケジュールを作成する今すぐ登録する

TPUを使用する

TensorFlow.orgで表示 GoogleColabで実行 GitHubでソースを表示ノートブックをダウンロードする

Cloud TPUの実験的なサポートは、現在KerasとGoogleColabで利用できます。このColabノートブックを実行する前に、ノートブックの設定を確認して、ハードウェアアクセラレータがTPUであることを確認してください:ランタイム>ランタイムタイプの変更>ハードウェアアクセラレータ> TPU

セットアップ

import tensorflow as tf

import os
import tensorflow_datasets as tfds

TPUの初期化

TPUは通常、クラウドTPUワーカーであり、ユーザーのPythonプログラムを実行するローカルプロセスとは異なります。したがって、リモートクラスターに接続してTPUを初期化するには、初期化作業を行う必要があります。 tpuへのtpu引数は、 tf.distribute.cluster_resolver.TPUClusterResolver専用の特別なアドレスであることに注意してください。 Google Compute Engine(GCE)でコードを実行している場合は、代わりにCloudTPUの名前を渡す必要があります。

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
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.42:8470
INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.42: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:3', 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')]

デバイスの手動配置

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でモデルを実行します。モデルを複数のTPU(または他のアクセラレータ)に配布するために、TensorFlowはいくつかの配布戦略を提供します。配布戦略を置き換えることができ、モデルは任意の(TPU)デバイスで実行されます。詳細については、配布戦略ガイドを確認してください。

これを示すために、 tf.distribute.TPUStrategyオブジェクトを作成します。

strategy = tf.distribute.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: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渡すことができます。以下は、すべてのコアが同じ入力(a, b)を受け取り、各コアで独立して行列乗算を実行(a, b)ことを示す例です。出力は、すべてのレプリカからの値になります。

@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.TPUStrategyを使用してクラウドTPUでtf.distribute.TPUStrategyモデルをトレーニングする方法を示します。

Kerasモデルを定義する

Kerasを使用したMNISTデータセットでの画像分類のためのSequentialモデルの定義から始めます。 CPUまたはGPUでトレーニングしている場合に使用するものと同じです。 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 APIを効率的に使用することは、Cloud TPUを使用する場合に重要です。これは、tf.data.Datasettf.data.Datasetデータをフィードできない限り、tf.data.Datasetを使用できないためです。データセットのパフォーマンスについて詳しくは、入力パイプラインパフォーマンスガイドをご覧ください

最も単純な実験( tf.data.Dataset.from_tensor_slicesまたはその他のグラフ内データを使用)を除くすべての場合、データセットによって読み取られたすべてのデータファイルをGoogle Cloud Storage(GCS)バケットに保存する必要があります。

ほとんどのユースケースでは、あなたにデータを変換することをお勧めしTFRecord形式と使用tf.data.TFRecordDatasetそれを読みます。これを行う方法の詳細については、 TFRecordとtf.Exampleチュートリアルを確認してください。これは難しい要件ではなく、 tf.data.FixedLengthRecordDatasettf.data.TextLineDatasetなどの他のデータセットリーダーを使用できます。

tf.data.Dataset.cacheを使用して、小さなデータセット全体をメモリにロードできます。

使用するデータ形式に関係なく、100MB程度の大きなファイルを使用することを強くお勧めします。これは、ファイルを開くオーバーヘッドが大幅に高くなるため、このネットワーク設定では特に重要です。

以下のコードに示すように、 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)

  # Normalize the input data.
  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 of having an
  # infinite dataset for training is to avoid the potential last partial batch
  # in each epoch, so that you 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のを使用してモデルを訓練することができますfitしてcompileするAPI。このステップにはTPU固有のものはありませんTPUStrategy代わりに複数のGPUとMirroredStrategyを使用しているかのようにコードを記述します。詳細については、 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
validation_steps = 10000 // 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, 
          validation_steps=validation_steps)
Epoch 1/5
300/300 [==============================] - 18s 32ms/step - loss: 0.1244 - sparse_categorical_accuracy: 0.9611 - val_loss: 0.0481 - val_sparse_categorical_accuracy: 0.9840
Epoch 2/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0322 - sparse_categorical_accuracy: 0.9898 - val_loss: 0.0375 - val_sparse_categorical_accuracy: 0.9875
Epoch 3/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0189 - sparse_categorical_accuracy: 0.9939 - val_loss: 0.0472 - val_sparse_categorical_accuracy: 0.9865
Epoch 4/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0117 - sparse_categorical_accuracy: 0.9960 - val_loss: 0.0474 - val_sparse_categorical_accuracy: 0.9874
Epoch 5/5
300/300 [==============================] - 7s 22ms/step - loss: 0.0084 - sparse_categorical_accuracy: 0.9973 - val_loss: 0.0431 - val_sparse_categorical_accuracy: 0.9891
<tensorflow.python.keras.callbacks.History at 0x7fa5ec4b9518>

Pythonのオーバーヘッドを減らし、TPUのパフォーマンスを最大化するには、実験的な引数( experimental_steps_per_executionModel.compileます。この例では、スループットが約50%増加します。

with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                # Anything between 2 and `steps_per_epoch` could help here.
                experimental_steps_per_execution = 50,
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['sparse_categorical_accuracy'])

model.fit(train_dataset,
          epochs=5,
          steps_per_epoch=steps_per_epoch,
          validation_data=test_dataset,
          validation_steps=validation_steps)
WARNING:tensorflow:The argument `steps_per_execution` is no longer experimental. Pass `steps_per_execution` instead of `experimental_steps_per_execution`.
WARNING:tensorflow:The argument `steps_per_execution` is no longer experimental. Pass `steps_per_execution` instead of `experimental_steps_per_execution`.
Epoch 1/5
300/300 [==============================] - 12s 39ms/step - loss: 0.1411 - sparse_categorical_accuracy: 0.9558 - val_loss: 0.0446 - val_sparse_categorical_accuracy: 0.9848
Epoch 2/5
300/300 [==============================] - 3s 9ms/step - loss: 0.0343 - sparse_categorical_accuracy: 0.9895 - val_loss: 0.0356 - val_sparse_categorical_accuracy: 0.9891
Epoch 3/5
300/300 [==============================] - 3s 9ms/step - loss: 0.0178 - sparse_categorical_accuracy: 0.9949 - val_loss: 0.0525 - val_sparse_categorical_accuracy: 0.9856
Epoch 4/5
300/300 [==============================] - 3s 9ms/step - loss: 0.0133 - sparse_categorical_accuracy: 0.9956 - val_loss: 0.0410 - val_sparse_categorical_accuracy: 0.9891
Epoch 5/5
300/300 [==============================] - 3s 9ms/step - loss: 0.0109 - sparse_categorical_accuracy: 0.9962 - val_loss: 0.0461 - val_sparse_categorical_accuracy: 0.9874
<tensorflow.python.keras.callbacks.History at 0x7fa5841d36a0>

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

tf.functionおよびtf.distribute直接使用して、モデルを作成およびトレーニングすることもできます。 strategy.experimental_distribute_datasets_from_function APIを使用して、データセット関数を指定してデータセットを配布できます。以下の例では、データセットに渡されるバッチサイズは、グローバルバッチサイズではなく、レプリカごとのバッチサイズであることに注意してください。詳細については、 tf.distribute.Strategyチュートリアルを使用したカスタムトレーニングをご覧ください。

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

# Create the model, optimizer and metrics inside the 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),))
WARNING:tensorflow:From <ipython-input-1-5625c2a14441>:15: StrategyBase.experimental_distribute_datasets_from_function (from tensorflow.python.distribute.distribute_lib) is deprecated and will be removed in a future version.
Instructions for updating:
rename to distribute_datasets_from_function
WARNING:tensorflow:From <ipython-input-1-5625c2a14441>:15: StrategyBase.experimental_distribute_datasets_from_function (from tensorflow.python.distribute.distribute_lib) is deprecated and will be removed in a future version.
Instructions for updating:
rename to distribute_datasets_from_function

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

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.1245, accuracy: 96.17%
Epoch: 1/5
Current step: 600, training loss: 0.0324, accuracy: 99.0%
Epoch: 2/5
Current step: 900, training loss: 0.0206, accuracy: 99.34%
Epoch: 3/5
Current step: 1200, training loss: 0.0128, accuracy: 99.56%
Epoch: 4/5
Current step: 1500, training loss: 0.0105, accuracy: 99.66%

tf.function内の複数のステップによるパフォーマンスの向上

tf.function内で複数のステップを実行することにより、パフォーマンスを向上させることができます。これは、 strategy.run呼び出しをtf.range内のtf.functionでラップすることによって実現され、AutoGraphはそれをTPUワーカーでtf.while_loopに変換します。

パフォーマンスは向上していますが、 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.0067, accuracy: 99.77%

次のステップ