Google I / O kehrt vom 18. bis 20. Mai zurück! Reservieren Sie Platz und erstellen Sie Ihren Zeitplan Registrieren Sie sich jetzt
Diese Seite wurde von der Cloud Translation API übersetzt.
Switch to English

Verteiltes Training mit Keras

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

Überblick

Die API tf.distribute.Strategy bietet eine Abstraktion für die Verteilung Ihres Trainings auf mehrere Verarbeitungseinheiten. Ziel ist es, Benutzern zu ermöglichen, verteiltes Training mit vorhandenen Modellen und Trainingscode mit minimalen Änderungen zu ermöglichen.

In diesem Lernprogramm wird die tf.distribute.MirroredStrategy , die die In-Graph-Replikation mit synchronem Training auf vielen GPUs auf einem Computer tf.distribute.MirroredStrategy . Im Wesentlichen werden alle Variablen des Modells auf jeden Prozessor kopiert. Anschließend werden die Farbverläufe aller Prozessoren mit all-redu kombiniert und der kombinierte Wert auf alle Kopien des Modells angewendet.

MirroredStrategy ist eine von mehreren Vertriebsstrategien, die im TensorFlow-Kern verfügbar sind. Weitere Strategien finden Sie im Distributionsstrategie-Handbuch .

Keras API

In diesem Beispiel wird die tf.keras API verwendet, um das Modell und die Trainingsschleife zu erstellen. Informationen zu benutzerdefinierten Trainingsschleifen finden Sie im Tutorial tf.distribute.Strategy with training loops .

Abhängigkeiten importieren

# Import TensorFlow and TensorFlow Datasets

import tensorflow_datasets as tfds
import tensorflow as tf

import os
print(tf.__version__)
2.4.1

Laden Sie den Datensatz herunter

Laden Sie das MNIST-Dataset herunter und laden Sie es aus den TensorFlow-Datasets . Dies gibt einen Datensatz im tf.data Format zurück.

Wenn Sie with_info auf True die Metadaten für das gesamte Dataset eingeschlossen, die hier in info gespeichert werden. Dieses Metadatenobjekt enthält unter anderem die Anzahl der Zug- und Testbeispiele.

datasets, info = tfds.load(name='mnist', with_info=True, as_supervised=True)

mnist_train, mnist_test = datasets['train'], datasets['test']

Vertriebsstrategie definieren

Erstellen Sie ein MirroredStrategy Objekt. Dies übernimmt die Verteilung und bietet einen Kontextmanager ( tf.distribute.MirroredStrategy.scope ), in dem Sie Ihr Modell erstellen können.

strategy = tf.distribute.MirroredStrategy()
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)
print('Number of devices: {}'.format(strategy.num_replicas_in_sync))
Number of devices: 1

Eingabepipeline einrichten

Wenn Sie ein Modell mit mehreren GPUs trainieren, können Sie die zusätzliche Rechenleistung effektiv nutzen, indem Sie die Stapelgröße erhöhen. Verwenden Sie im Allgemeinen die größte Stapelgröße, die zum GPU-Speicher passt, und passen Sie die Lernrate entsprechend an.

# You can also do info.splits.total_num_examples to get the total
# number of examples in the dataset.

num_train_examples = info.splits['train'].num_examples
num_test_examples = info.splits['test'].num_examples

BUFFER_SIZE = 10000

BATCH_SIZE_PER_REPLICA = 64
BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync

Pixelwerte von 0 bis 255 müssen auf den Bereich 0 bis 1 normiert werden . Definieren Sie diese Skala in einer Funktion.

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

  return image, label

Wenden Sie diese Funktion auf die Trainings- und Testdaten an, mischen Sie die Trainingsdaten und stapeln Sie sie für das Training . Beachten Sie, dass wir auch einen In-Memory-Cache der Trainingsdaten aufbewahren, um die Leistung zu verbessern.

train_dataset = mnist_train.map(scale).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
eval_dataset = mnist_test.map(scale).batch(BATCH_SIZE)

Erstellen Sie das Modell

Erstellen und kompilieren Sie das Keras-Modell im Kontext von strategy.scope .

with strategy.scope():
  model = tf.keras.Sequential([
      tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),
      tf.keras.layers.MaxPooling2D(),
      tf.keras.layers.Flatten(),
      tf.keras.layers.Dense(64, activation='relu'),
      tf.keras.layers.Dense(10)
  ])

  model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                optimizer=tf.keras.optimizers.Adam(),
                metrics=['accuracy'])
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).

Definieren Sie die Rückrufe

Die hier verwendeten Rückrufe sind:

  • TensorBoard : Dieser Rückruf schreibt ein Protokoll für TensorBoard, mit dem Sie die Diagramme visualisieren können.
  • Modellprüfpunkt : Dieser Rückruf speichert das Modell nach jeder Epoche.
  • Scheduler für Lernraten: Mit diesem Rückruf können Sie festlegen, dass sich die Lernrate nach jeder Epoche / Charge ändert.

Fügen Sie zur Veranschaulichung einen Druckrückruf hinzu, um die Lernrate im Notizbuch anzuzeigen.

# Define the checkpoint directory to store the checkpoints

checkpoint_dir = './training_checkpoints'
# Name of the checkpoint files
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt_{epoch}")
# Function for decaying the learning rate.
# You can define any decay function you need.
def decay(epoch):
  if epoch < 3:
    return 1e-3
  elif epoch >= 3 and epoch < 7:
    return 1e-4
  else:
    return 1e-5
# Callback for printing the LR at the end of each epoch.
class PrintLR(tf.keras.callbacks.Callback):
  def on_epoch_end(self, epoch, logs=None):
    print('\nLearning rate for epoch {} is {}'.format(epoch + 1,
                                                      model.optimizer.lr.numpy()))
callbacks = [
    tf.keras.callbacks.TensorBoard(log_dir='./logs'),
    tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_prefix,
                                       save_weights_only=True),
    tf.keras.callbacks.LearningRateScheduler(decay),
    PrintLR()
]

Trainieren und bewerten

Trainieren Sie nun das Modell auf die übliche Weise, rufen Sie fit für das Modell auf und übergeben Sie den zu Beginn des Lernprogramms erstellten Datensatz. Dieser Schritt ist der gleiche, unabhängig davon, ob Sie das Training verteilen oder nicht.

model.fit(train_dataset, epochs=12, callbacks=callbacks)
Epoch 1/12
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
5/938 [..............................] - ETA: 11s - loss: 2.2256 - accuracy: 0.2527    WARNING:tensorflow:Callback method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0036s vs `on_train_batch_end` time: 0.0062s). Check your callbacks.
WARNING:tensorflow:Callback method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0036s vs `on_train_batch_end` time: 0.0062s). Check your callbacks.
938/938 [==============================] - 18s 4ms/step - loss: 0.3872 - accuracy: 0.8928

Learning rate for epoch 1 is 0.0010000000474974513
Epoch 2/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0713 - accuracy: 0.9787

Learning rate for epoch 2 is 0.0010000000474974513
Epoch 3/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0512 - accuracy: 0.9855

Learning rate for epoch 3 is 0.0010000000474974513
Epoch 4/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0276 - accuracy: 0.9919

Learning rate for epoch 4 is 9.999999747378752e-05
Epoch 5/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0251 - accuracy: 0.9930

Learning rate for epoch 5 is 9.999999747378752e-05
Epoch 6/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0225 - accuracy: 0.9940

Learning rate for epoch 6 is 9.999999747378752e-05
Epoch 7/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0204 - accuracy: 0.9948

Learning rate for epoch 7 is 9.999999747378752e-05
Epoch 8/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0165 - accuracy: 0.9959

Learning rate for epoch 8 is 9.999999747378752e-06
Epoch 9/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0180 - accuracy: 0.9955

Learning rate for epoch 9 is 9.999999747378752e-06
Epoch 10/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0187 - accuracy: 0.9955

Learning rate for epoch 10 is 9.999999747378752e-06
Epoch 11/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0177 - accuracy: 0.9959

Learning rate for epoch 11 is 9.999999747378752e-06
Epoch 12/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0170 - accuracy: 0.9961

Learning rate for epoch 12 is 9.999999747378752e-06
<tensorflow.python.keras.callbacks.History at 0x7fa6a5c1f9e8>

Wie Sie unten sehen können, werden die Kontrollpunkte gespeichert.

# check the checkpoint directory
ls {checkpoint_dir}
checkpoint           ckpt_4.data-00000-of-00001
ckpt_1.data-00000-of-00001   ckpt_4.index
ckpt_1.index             ckpt_5.data-00000-of-00001
ckpt_10.data-00000-of-00001  ckpt_5.index
ckpt_10.index            ckpt_6.data-00000-of-00001
ckpt_11.data-00000-of-00001  ckpt_6.index
ckpt_11.index            ckpt_7.data-00000-of-00001
ckpt_12.data-00000-of-00001  ckpt_7.index
ckpt_12.index            ckpt_8.data-00000-of-00001
ckpt_2.data-00000-of-00001   ckpt_8.index
ckpt_2.index             ckpt_9.data-00000-of-00001
ckpt_3.data-00000-of-00001   ckpt_9.index
ckpt_3.index

Um zu sehen , wie das Modell ausführen, laden Sie die neueste Kontrollpunkt und Anruf zu evaluate auf der Testdaten.

Rufen Sie evaluate wie zuvor unter Verwendung geeigneter Datensätze auf.

model.load_weights(tf.train.latest_checkpoint(checkpoint_dir))

eval_loss, eval_acc = model.evaluate(eval_dataset)

print('Eval loss: {}, Eval Accuracy: {}'.format(eval_loss, eval_acc))
157/157 [==============================] - 3s 5ms/step - loss: 0.0382 - accuracy: 0.9871
Eval loss: 0.03821386396884918, Eval Accuracy: 0.9871000051498413

Um die Ausgabe anzuzeigen, können Sie die TensorBoard-Protokolle am Terminal herunterladen und anzeigen.

$ tensorboard --logdir=path/to/log-directory
ls -sh ./logs
total 4.0K
4.0K train

Export nach SavedModel

Exportieren Sie das Diagramm und die Variablen in das plattformunabhängige SavedModel-Format. Nachdem Ihr Modell gespeichert wurde, können Sie es mit oder ohne Gültigkeitsbereich laden.

path = 'saved_model/'
model.save(path, save_format='tf')
INFO:tensorflow:Assets written to: saved_model/assets
INFO:tensorflow:Assets written to: saved_model/assets

Laden Sie das Modell ohne strategy.scope .

unreplicated_model = tf.keras.models.load_model(path)

unreplicated_model.compile(
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=tf.keras.optimizers.Adam(),
    metrics=['accuracy'])

eval_loss, eval_acc = unreplicated_model.evaluate(eval_dataset)

print('Eval loss: {}, Eval Accuracy: {}'.format(eval_loss, eval_acc))
157/157 [==============================] - 1s 3ms/step - loss: 0.0364 - accuracy: 0.9883
Eval loss: 0.03821386396884918, Eval Accuracy: 0.9871000051498413

Laden Sie das Modell mit strategy.scope .

with strategy.scope():
  replicated_model = tf.keras.models.load_model(path)
  replicated_model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                           optimizer=tf.keras.optimizers.Adam(),
                           metrics=['accuracy'])

  eval_loss, eval_acc = replicated_model.evaluate(eval_dataset)
  print ('Eval loss: {}, Eval Accuracy: {}'.format(eval_loss, eval_acc))
157/157 [==============================] - 3s 4ms/step - loss: 0.0364 - accuracy: 0.9883
Eval loss: 0.03821386396884918, Eval Accuracy: 0.9871000051498413

Beispiele und Tutorials

Hier einige Beispiele für die Verwendung der Verteilungsstrategie mit Keras Fit / Compile:

  1. Beispiel für einen Transformator, der mit tf.distribute.MirroredStrategy trainiert wurde
  2. NCF- Beispiel, das mit tf.distribute.MirroredStrategy trainiert wurde.

Weitere Beispiele finden Sie im Handbuch zur Vertriebsstrategie

Nächste Schritte