Cette page a été traduite par l'API Cloud Translation.
Switch to English

Formation distribuée avec Keras

Voir sur TensorFlow.org Exécuter dans Google Colab Afficher la source sur GitHub Télécharger le carnet

Aperçu

L'API tf.distribute.Strategy fournit une abstraction pour distribuer votre formation sur plusieurs unités de traitement. L'objectif est de permettre aux utilisateurs d'activer la formation distribuée à l'aide des modèles et du code de formation existants, avec des modifications minimes.

Ce didacticiel utilise tf.distribute.MirroredStrategy , qui effectue une réplication dans le graphe avec entraînement synchrone sur de nombreux GPU sur une seule machine. Essentiellement, il copie toutes les variables du modèle dans chaque processeur. Ensuite, il utilise all-reduction pour combiner les dégradés de tous les processeurs et applique la valeur combinée à toutes les copies du modèle.

MirroredStrategy est l'une des nombreuses stratégies de distribution disponibles dans TensorFlow core. Vous pouvez en savoir plus sur plus de stratégies dans le guide de stratégie de distribution .

API Keras

Cet exemple utilise l'API tf.keras pour créer le modèle et la boucle d'entraînement. Pour les boucles d'entraînement personnalisées, consultez le didacticiel tf.distribute.Strategy avec boucles d'entraînement .

Importer les dépendances

# Import TensorFlow and TensorFlow Datasets

import tensorflow_datasets as tfds
import tensorflow as tf
tfds.disable_progress_bar()

import os
print(tf.__version__)
2.2.0

Téléchargez le jeu de données

Téléchargez le jeu de données MNIST et chargez-le à partir des jeux de données TensorFlow . Cela renvoie un ensemble de données au format tf.data .

La définition de with_info sur True inclut les métadonnées pour l'ensemble de données, qui sont enregistrées ici dans info . Entre autres choses, cet objet de métadonnées inclut le nombre d'exemples de train et de test.

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

mnist_train, mnist_test = datasets['train'], datasets['test']
Downloading and preparing dataset mnist/3.0.1 (download: 11.06 MiB, generated: 21.00 MiB, total: 32.06 MiB) to /home/kbuilder/tensorflow_datasets/mnist/3.0.1...

Warning:absl:Dataset mnist is hosted on GCS. It will automatically be downloaded to your
local data directory. If you'd instead prefer to read directly from our public
GCS bucket (recommended if you're running on GCP), you can instead pass
`try_gcs=True` to `tfds.load` or set `data_dir=gs://tfds-data/datasets`.


Dataset mnist downloaded and prepared to /home/kbuilder/tensorflow_datasets/mnist/3.0.1. Subsequent calls will reuse this data.

Définir la stratégie de distribution

Créez un objet MirroredStrategy . Cela gérera la distribution et fournira un gestionnaire de contexte ( tf.distribute.MirroredStrategy.scope ) pour construire votre modèle à l'intérieur.

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

Configurer le pipeline d'entrée

Lors de la formation d'un modèle avec plusieurs GPU, vous pouvez utiliser efficacement la puissance de calcul supplémentaire en augmentant la taille du lot. En général, utilisez la plus grande taille de lot adaptée à la mémoire du GPU et ajustez le taux d'apprentissage en conséquence.

# 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

Les valeurs de pixel, qui vont de 0 à 255, doivent être normalisées dans la plage 0-1 . Définissez cette échelle dans une fonction.

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

  return image, label

Appliquez cette fonction aux données d'entraînement et de test, mélangez les données d'entraînement et regroupez-les pour l'entraînement . Notez que nous conservons également un cache en mémoire des données d'entraînement pour améliorer les performances.

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

Créer le modèle

Créez et compilez le modèle Keras dans le contexte de 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'])

Définir les rappels

Les rappels utilisés ici sont:

  • TensorBoard : Ce rappel écrit un journal pour TensorBoard qui vous permet de visualiser les graphiques.
  • Point de contrôle du modèle : ce rappel enregistre le modèle après chaque époque.
  • Programmateur de taux d'apprentissage : à l'aide de ce rappel, vous pouvez programmer le taux d'apprentissage pour qu'il change après chaque époque / lot.

À des fins d'illustration, ajoutez un rappel d'impression pour afficher le taux d'apprentissage dans le bloc-notes.

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

Former et évaluer

Maintenant, entraînez le modèle de la manière habituelle, en appelant fit sur le modèle et en passant le jeu de données créé au début du didacticiel. Cette étape est la même que vous distribuiez la formation ou non.

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',).

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',).

936/938 [============================>.] - ETA: 0s - accuracy: 0.9422 - loss: 0.2016
Learning rate for epoch 1 is 0.0010000000474974513
938/938 [==============================] - 5s 5ms/step - accuracy: 0.9422 - loss: 0.2015 - lr: 0.0010
Epoch 2/12
936/938 [============================>.] - ETA: 0s - accuracy: 0.9807 - loss: 0.0662
Learning rate for epoch 2 is 0.0010000000474974513
938/938 [==============================] - 3s 4ms/step - accuracy: 0.9807 - loss: 0.0662 - lr: 0.0010
Epoch 3/12
933/938 [============================>.] - ETA: 0s - accuracy: 0.9863 - loss: 0.0464
Learning rate for epoch 3 is 0.0010000000474974513
938/938 [==============================] - 3s 3ms/step - accuracy: 0.9863 - loss: 0.0464 - lr: 0.0010
Epoch 4/12
933/938 [============================>.] - ETA: 0s - accuracy: 0.9933 - loss: 0.0252
Learning rate for epoch 4 is 9.999999747378752e-05
938/938 [==============================] - 3s 4ms/step - accuracy: 0.9933 - loss: 0.0252 - lr: 1.0000e-04
Epoch 5/12
932/938 [============================>.] - ETA: 0s - accuracy: 0.9946 - loss: 0.0220
Learning rate for epoch 5 is 9.999999747378752e-05
938/938 [==============================] - 3s 3ms/step - accuracy: 0.9945 - loss: 0.0220 - lr: 1.0000e-04
Epoch 6/12
929/938 [============================>.] - ETA: 0s - accuracy: 0.9951 - loss: 0.0200
Learning rate for epoch 6 is 9.999999747378752e-05
938/938 [==============================] - 3s 3ms/step - accuracy: 0.9951 - loss: 0.0201 - lr: 1.0000e-04
Epoch 7/12
928/938 [============================>.] - ETA: 0s - accuracy: 0.9955 - loss: 0.0186
Learning rate for epoch 7 is 9.999999747378752e-05
938/938 [==============================] - 3s 3ms/step - accuracy: 0.9955 - loss: 0.0186 - lr: 1.0000e-04
Epoch 8/12
934/938 [============================>.] - ETA: 0s - accuracy: 0.9965 - loss: 0.0161
Learning rate for epoch 8 is 9.999999747378752e-06
938/938 [==============================] - 3s 3ms/step - accuracy: 0.9965 - loss: 0.0161 - lr: 1.0000e-05
Epoch 9/12
932/938 [============================>.] - ETA: 0s - accuracy: 0.9965 - loss: 0.0157
Learning rate for epoch 9 is 9.999999747378752e-06
938/938 [==============================] - 3s 3ms/step - accuracy: 0.9965 - loss: 0.0156 - lr: 1.0000e-05
Epoch 10/12
934/938 [============================>.] - ETA: 0s - accuracy: 0.9966 - loss: 0.0155
Learning rate for epoch 10 is 9.999999747378752e-06
938/938 [==============================] - 3s 3ms/step - accuracy: 0.9966 - loss: 0.0154 - lr: 1.0000e-05
Epoch 11/12
934/938 [============================>.] - ETA: 0s - accuracy: 0.9967 - loss: 0.0153
Learning rate for epoch 11 is 9.999999747378752e-06
938/938 [==============================] - 3s 3ms/step - accuracy: 0.9967 - loss: 0.0153 - lr: 1.0000e-05
Epoch 12/12
924/938 [============================>.] - ETA: 0s - accuracy: 0.9967 - loss: 0.0152
Learning rate for epoch 12 is 9.999999747378752e-06
938/938 [==============================] - 3s 3ms/step - accuracy: 0.9967 - loss: 0.0151 - lr: 1.0000e-05

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

Comme vous pouvez le voir ci-dessous, les points de contrôle sont enregistrés.

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

Pour voir comment le modèle fonctionne, chargez le dernier point de contrôle et appelez evaluate sur les données de test.

Appelez evaluate comme avant d'utiliser les ensembles de données appropriés.

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))
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',).

157/157 [==============================] - 1s 7ms/step - accuracy: 0.9861 - loss: 0.0393
Eval loss: 0.039307601749897, Eval Accuracy: 0.9861000180244446

Pour voir la sortie, vous pouvez télécharger et afficher les journaux TensorBoard sur le terminal.

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

Exporter vers SavedModel

Exportez le graphique et les variables au format SavedModel indépendant de la plate-forme. Une fois votre modèle enregistré, vous pouvez le charger avec ou sans la portée.

path = 'saved_model/'
model.save(path, save_format='tf')
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/ops/resource_variable_ops.py:1817: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.
Instructions for updating:
If using Keras pass *_constraint arguments to layers.

Warning:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/ops/resource_variable_ops.py:1817: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.
Instructions for updating:
If using Keras pass *_constraint arguments to layers.

INFO:tensorflow:Assets written to: saved_model/assets

INFO:tensorflow:Assets written to: saved_model/assets

Chargez le modèle sans 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 5ms/step - loss: 0.0393 - accuracy: 0.9861
Eval loss: 0.039307601749897, Eval Accuracy: 0.9861000180244446

Chargez le modèle avec 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 [==============================] - 1s 6ms/step - accuracy: 0.9861 - loss: 0.0393
Eval loss: 0.039307601749897, Eval Accuracy: 0.9861000180244446

Exemples et tutoriels

Voici quelques exemples d'utilisation de la stratégie de distribution avec keras fit / compile:

  1. Exemple de transformateur formé à l'aide de tf.distribute.MirroredStrategy
  2. Exemple NCF formé à l'aide de tf.distribute.MirroredStrategy .

Plus d'exemples répertoriés dans le guide de stratégie de distribution

Prochaines étapes