TPUs verwenden

Auf TensorFlow.org ansehen In Google Colab ausführen Quelle auf GitHub anzeigen Notizbuch herunterladen

Bevor Sie dieses Colab Notebook ausführen, stellen Sie sicher , dass Ihre Hardware - Beschleuniger durch Überprüfen Ihres Notebook - Einstellungen ein TPU ist: Runtime> Ändern Laufzeittyp> Hardware - Beschleuniger> TPU.

Installieren

import tensorflow as tf

import os
import tensorflow_datasets as tfds
/home/kbuilder/.local/lib/python3.6/site-packages/requests/__init__.py:104: RequestsDependencyWarning: urllib3 (1.26.6) or chardet (2.3.0)/charset_normalizer (2.0.4) doesn't match a supported version!
  RequestsDependencyWarning)

TPU-Initialisierung

TPUs sind in der Regel Cloud TPU-Worker, die sich von dem lokalen Prozess unterscheiden, der das Python-Programm des Nutzers ausführt. Daher müssen Sie einige Initialisierungsarbeiten durchführen, um eine Verbindung zum Remote-Cluster herzustellen und die TPUs zu initialisieren. Beachten Sie, dass das tpu Argument tf.distribute.cluster_resolver.TPUClusterResolver eine spezielle Adresse ist nur für Colab. Wenn Sie Ihren Code auf Google Compute Engine (GCE) ausführen, sollten Sie stattdessen den Namen Ihrer Cloud TPU übergeben.

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:Clearing out eager caches
INFO:tensorflow:Clearing out eager caches
INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.2:8470
INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.2:8470
INFO:tensorflow:Finished initializing TPU system.
INFO:tensorflow:Finished initializing TPU system.
All devices:  [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'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:4', 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:6', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:7', device_type='TPU')]

Manuelle Geräteplatzierung

Nachdem die TPU initialisiert wurde, können Sie die manuelle Geräteplatzierung verwenden, um die Berechnung auf einem einzelnen TPU-Gerät zu platzieren:

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)

Vertriebsstrategien

Normalerweise führen Sie Ihr Modell auf mehreren TPUs datenparallel aus. Um Ihr Modell auf mehreren TPUs (oder anderen Beschleunigern) zu verteilen, bietet TensorFlow mehrere Verteilungsstrategien. Sie können Ihre Verteilungsstrategie ersetzen und das Modell wird auf jedem beliebigen (TPU) Gerät ausgeführt. Schauen Sie sich die Vertriebsstrategie Leitfaden für weitere Informationen.

Um dies zu demonstrieren, erstellen Sie ein tf.distribute.TPUStrategy Objekt:

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)

Um eine Berechnung zu replizieren , so dass es in allen TPU Kerne laufen können, können Sie es in den Pass strategy.run API. Unten ist ein Beispiel , das zeigt alle Kerne die gleichen Eingänge empfangen (a, b) auf jedem Kern unabhängig und Durchführen einer Matrixmultiplikation. Die Ausgaben sind die Werte aller Replikate.

@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)
}

Klassifizierung nach TPUs

Nachdem Sie die grundlegenden Konzepte behandelt haben, betrachten Sie ein konkreteres Beispiel. In diesem Abschnitt wird gezeigt , wie die Strategie- Verteilung verwendet werden tf.distribute.TPUStrategy -eine Keras Modell auf einer Wolke TPU zu trainieren.

Definieren Sie ein Keras-Modell

Beginnen Sie mit einer Definition eines Sequential Keras Modell für Bildklassifikation auf dem MNIST Dataset Keras. Es unterscheidet sich nicht von dem, was Sie verwenden würden, wenn Sie auf CPUs oder GPUs trainieren. Beachten Sie, dass Keras Modell Bedürfnisse Schaffung innen sein strategy.scope , so können die Variablen für jedes TPU - Gerät erstellt werden. Andere Teile des Codes sind nicht erforderlich, um innerhalb des Strategiebereichs zu liegen.

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)])

Laden Sie den Datensatz

Effiziente Nutzung der tf.data.Dataset API ist kritisch , wenn eine Wolke TPU verwenden, da es unmöglich ist , den Cloud - TPUs zu verwenden , wenn Sie schnell genug , um sie Daten füttern können. Sie können mehr über Daten - Set Leistung in der Lerneingang Pipeline Performance Führung .

Für alle , aber die einfachsten Experimente (mit tf.data.Dataset.from_tensor_slices oder andere in-Graph - Daten), müssen Sie alle Dateien , die vom Datensatz in Google Cloud Storage (GCS) Eimer lesen speichern.

Für die meisten Anwendungsfälle ist es empfehlenswert , Ihre Daten in das konvertieren TFRecord Format und verwenden Sie einen tf.data.TFRecordDataset es zu lesen. Überprüfen Sie die TFRecord und tf.Example Tutorial für Details, wie dies zu tun. Es ist nicht eine harte Anforderung und Sie können andere Datenmenge Leser, wie verwenden tf.data.FixedLengthRecordDataset oder tf.data.TextLineDataset .

Sie können ganze kleine Datensätze in den Speicher unter Verwendung laden tf.data.Dataset.cache .

Unabhängig vom verwendeten Datenformat wird dringend empfohlen, große Dateien in der Größenordnung von 100 MB zu verwenden. Dies ist in dieser vernetzten Umgebung besonders wichtig, da der Aufwand beim Öffnen einer Datei deutlich höher ist.

Wie weiter unten in dem Code angezeigt, sollten Sie die Verwendung tensorflow_datasets Modul eine Kopie der MNIST Trainings- und Testdaten zu erhalten. Beachten Sie, dass try_gcs angegeben ist , eine Kopie zu verwenden , die in einem öffentlichen GCS Eimer zur Verfügung steht. Wenn Sie dies nicht angeben, kann die TPU nicht auf die heruntergeladenen Daten zugreifen.

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

Trainieren Sie das Modell mit den High-Level-APIs von Keras

Sie können Ihr Modell mit Keras trainieren fit und compile APIs. Es gibt nichts , TPU-spezifische in diesem Schritt-Sie den Code schreiben , als ob Sie mutliple GPUs und eine Verwendung wurden MirroredStrategy anstelle des TPUStrategy . Mehr erfahren Sie in dem lernen mit Keras Distributed Training Tutorial.

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 35ms/step - loss: 0.1352 - sparse_categorical_accuracy: 0.9583 - val_loss: 0.0404 - val_sparse_categorical_accuracy: 0.9863
Epoch 2/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0346 - sparse_categorical_accuracy: 0.9893 - val_loss: 0.0376 - val_sparse_categorical_accuracy: 0.9882
Epoch 3/5
300/300 [==============================] - 7s 22ms/step - loss: 0.0199 - sparse_categorical_accuracy: 0.9934 - val_loss: 0.0432 - val_sparse_categorical_accuracy: 0.9872
Epoch 4/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0132 - sparse_categorical_accuracy: 0.9957 - val_loss: 0.0475 - val_sparse_categorical_accuracy: 0.9870
Epoch 5/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0100 - sparse_categorical_accuracy: 0.9970 - val_loss: 0.0443 - val_sparse_categorical_accuracy: 0.9877
<keras.callbacks.History at 0x7faf8829f278>

Python Aufwand zu reduzieren und die Leistung Ihres TPU zu maximieren, passieren in der argument- steps_per_execution -bis Model.compile . In diesem Beispiel wird der Durchsatz um etwa 50 % erhöht:

with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                # Anything between 2 and `steps_per_epoch` could help here.
                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)
Epoch 1/5
300/300 [==============================] - 13s 43ms/step - loss: 0.1405 - sparse_categorical_accuracy: 0.9571 - val_loss: 0.0438 - val_sparse_categorical_accuracy: 0.9860
Epoch 2/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0337 - sparse_categorical_accuracy: 0.9890 - val_loss: 0.0503 - val_sparse_categorical_accuracy: 0.9847
Epoch 3/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0170 - sparse_categorical_accuracy: 0.9947 - val_loss: 0.0448 - val_sparse_categorical_accuracy: 0.9884
Epoch 4/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0137 - sparse_categorical_accuracy: 0.9956 - val_loss: 0.0362 - val_sparse_categorical_accuracy: 0.9903
Epoch 5/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0090 - sparse_categorical_accuracy: 0.9971 - val_loss: 0.0458 - val_sparse_categorical_accuracy: 0.9880
<keras.callbacks.History at 0x7fae1054f908>

Trainieren Sie das Modell mit einer benutzerdefinierten Trainingsschleife

Sie können auch Ihre Modelle erstellen und trainieren tf.function und tf.distribute APIs direkt. Sie können die Verwendung strategy.experimental_distribute_datasets_from_function API den Datensatz gegeben ein Datensatz Funktion zu verteilen. Beachten Sie, dass im folgenden Beispiel die an das Dataset übergebene Batchgröße die Batchgröße pro Replikat und nicht die globale Batchgröße ist. Um mehr zu erfahren, besuchen Sie das Exklusiv - Training mit tf.distribute.Strategy Tutorial.

Erstellen Sie zunächst das Modell, die Datensätze und die 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

Führen Sie dann die Trainingsschleife aus:

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.1352, accuracy: 95.85%
Epoch: 1/5
Current step: 600, training loss: 0.0355, accuracy: 98.86%
Epoch: 2/5
Current step: 900, training loss: 0.0194, accuracy: 99.37%
Epoch: 3/5
Current step: 1200, training loss: 0.0127, accuracy: 99.6%
Epoch: 4/5
Current step: 1500, training loss: 0.0111, accuracy: 99.65%

Die Verbesserung der Leistung mit mehreren Schritten innerhalb tf.function

Sie können durch das Ausführen mehrerer Schritte innerhalb einer Verbesserung der Leistung tf.function . Dies wird durch Umwickeln des erreicht strategy.run Anruf mit einem tf.range innerhalb tf.function und Autograph wird es zu einem konvertieren tf.while_loop auf dem TPU Arbeiter.

Trotz der verbesserten Leistung, gibt es Nachteile bei diesem Verfahren im Vergleich zu einem einzigen Schritt innen verlauf tf.function . Mehrere Schritte in einem laufenden tf.function ist weniger flexibel : Sie können nicht laufen Dinge eifrig oder beliebigen Python - Code in den Schritten.

@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.0081, accuracy: 99.73%

Nächste Schritte