 TensorFlow.orgで表示 |  GoogleColabで実行 |  GitHubでソースを表示 |  ノートブックをダウンロード |
Cloud TPUの実験的なサポートは、現在KerasとGoogleColabで利用できます。このColabノートブックを実行する前に、ノートブックの設定を確認して、ハードウェアアクセラレータがTPUであることを確認してください：ランタイム>ランタイムタイプの変更>ハードウェアアクセラレータ> 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.18:8470 INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.18: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でモデルを実行したいと考えています。分散戦略は、CPU、GPU、またはTPUでモデルを駆動するために使用できる抽象化です。配布戦略を交換するだけで、モデルは特定のデバイスで実行されます。詳細については、流通戦略ガイドを参照してください。
まず、 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)を取得し、各コアで個別に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.TPUStrategyを使用してクラウドTPUを駆動し、 tf.distribute.TPUStrategyモデルをトレーニングする方法を示します。
Kerasモデルを定義する
以下は、Kerasを使用したMNISTモデルの定義であり、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を使用する場合に重要です。これは、Cloud TPUにデータを十分に迅速に提供できない限り、tf.data.Datasetを使用できないためです。データセットのパフォーマンスの詳細については、入力パイプラインパフォーマンスガイドを参照してください。
最も単純な実験（ tf.data.Dataset.from_tensor_slicesまたはその他のグラフ内データを使用）を除くすべての場合、データセットによって読み取られたすべてのデータファイルをGoogle Cloud Storage（GCS）バケットに保存する必要があります。
ほとんどのユースケースでは、データをTFRecord形式に変換し、 tf.data.TFRecordDatasetを使用して読み取ることをお勧めします。これを行う方法の詳細については、 TFRecordおよびtf.Exampleチュートリアルを参照してください。ただし、これは難しい要件ではなく、必要に応じて他のデータセットリーダー（ FixedLengthRecordDatasetまたは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)
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固有のものはありません。複数のGPUがあり、 TPUStrategyではなく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 [==============================] - 19s 37ms/step - loss: 0.3028 - sparse_categorical_accuracy: 0.9042 - val_loss: 0.0469 - val_sparse_categorical_accuracy: 0.9850 Epoch 2/5 300/300 [==============================] - 8s 25ms/step - loss: 0.0400 - sparse_categorical_accuracy: 0.9885 - val_loss: 0.0396 - val_sparse_categorical_accuracy: 0.9870 Epoch 3/5 300/300 [==============================] - 8s 25ms/step - loss: 0.0212 - sparse_categorical_accuracy: 0.9931 - val_loss: 0.0373 - val_sparse_categorical_accuracy: 0.9893 Epoch 4/5 300/300 [==============================] - 8s 25ms/step - loss: 0.0159 - sparse_categorical_accuracy: 0.9951 - val_loss: 0.0392 - val_sparse_categorical_accuracy: 0.9875 Epoch 5/5 300/300 [==============================] - 8s 25ms/step - loss: 0.0096 - sparse_categorical_accuracy: 0.9968 - val_loss: 0.0463 - val_sparse_categorical_accuracy: 0.9879 <tensorflow.python.keras.callbacks.History at 0x7f16c5c9e908>
Pythonのオーバーヘッドを減らし、TPUのパフォーマンスを最大化するには、 Model.compile実験experimental_steps_per_executionなexperimental_steps_per_execution引数をModel.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 [==============================] - 14s 45ms/step - loss: 0.2200 - sparse_categorical_accuracy: 0.9316 - val_loss: 0.0424 - val_sparse_categorical_accuracy: 0.9856 Epoch 2/5 300/300 [==============================] - 5s 15ms/step - loss: 0.0366 - sparse_categorical_accuracy: 0.9896 - val_loss: 0.0357 - val_sparse_categorical_accuracy: 0.9885 Epoch 3/5 300/300 [==============================] - 5s 15ms/step - loss: 0.0204 - sparse_categorical_accuracy: 0.9937 - val_loss: 0.0391 - val_sparse_categorical_accuracy: 0.9887 Epoch 4/5 300/300 [==============================] - 5s 15ms/step - loss: 0.0140 - sparse_categorical_accuracy: 0.9951 - val_loss: 0.0457 - val_sparse_categorical_accuracy: 0.9868 Epoch 5/5 300/300 [==============================] - 5s 16ms/step - loss: 0.0104 - sparse_categorical_accuracy: 0.9967 - val_loss: 0.0519 - val_sparse_categorical_accuracy: 0.9859 <tensorflow.python.keras.callbacks.History at 0x7f16c5d965c0>
カスタムトレーニングループを使用してモデルをトレーニングします。
tf.functionおよびtf.distribute直接使用して、モデルを作成およびトレーニングすることもできます。 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),))
WARNING:tensorflow:From <ipython-input-1-2f075cabff81>: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-2f075cabff81>: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.1313, accuracy: 95.91% Epoch: 1/5 Current step: 600, training loss: 0.0334, accuracy: 98.96% Epoch: 2/5 Current step: 900, training loss: 0.0189, accuracy: 99.42% Epoch: 3/5 Current step: 1200, training loss: 0.0131, accuracy: 99.57% Epoch: 4/5 Current step: 1500, training loss: 0.0104, accuracy: 99.64%
tf.function内の複数のステップによるパフォーマンスの向上
tf.function内で複数のステップを実行することにより、パフォーマンスを向上させることができます。これは、ラップすることによって達成されるstrategy.runでコールをtf.range内部tf.function 、オートグラフはに変換されますtf.while_loop TPUワーカーに。
パフォーマンスは向上しますが、 tf.function内の1つのステップと比較するとトレードオフがあります。 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.0087, accuracy: 99.73%
次のステップ
- Google CloudTPUのドキュメント-GoogleCloudTPUを設定して実行します。
- TensorFlowを使用した分散トレーニング-分散戦略の使用方法と、ベストプラクティスを示す多くの例へのリンク。
- TensorFlowを使用したモデルの保存/読み込み-分散戦略を使用してモデルを保存および読み込みする方法。
- TensorFlow公式モデル-クラウドTPUと互換性のある最先端のTensorFlow2.xモデルの例。
- Google CloudTPUパフォーマンスガイド-アプリケーションのCloudTPU構成パラメーターを調整して、CloudTPUのパフォーマンスをさらに向上させます。