Esta página foi traduzida pela API Cloud Translation.

Switch to English

Treinamento distribuído com Keras

Ver no TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar caderno

Visão geral

A API tf.distribute.Strategy fornece uma abstração para distribuir seu treinamento em várias unidades de processamento. O objetivo é permitir que os usuários habilitem o treinamento distribuído usando modelos existentes e código de treinamento, com mudanças mínimas.

Este tutorial usa o tf.distribute.MirroredStrategy , que faz replicação no gráfico com treinamento síncrono em muitas GPUs em uma máquina. Essencialmente, ele copia todas as variáveis do modelo para cada processador. Em seguida, ele usa all-reduzir para combinar os gradientes de todos os processadores e aplica o valor combinado a todas as cópias do modelo.

MirroredStrategy é uma das várias estratégias de distribuição disponíveis no núcleo do TensorFlow. Você pode ler sobre mais estratégias no guia de estratégia de distribuição .

API Keras

Este exemplo usa a API tf.keras para construir o modelo e o loop de treinamento. Para os loops de treinamento personalizados, consulte o tutorial tf.distribute.Strategy com loops de treinamento .

Importar dependências

# Import TensorFlow and TensorFlow Datasets

import tensorflow_datasets as tfds
import tensorflow as tf

import os

print(tf.__version__)

2.5.0

Baixe o conjunto de dados

Faça o download do conjunto de dados MNIST e carregue-o dos conjuntos de dados TensorFlow . Isso retorna um conjunto de dados no formato tf.data .

Definir with_info como True inclui os metadados para todo o conjunto de dados, que está sendo salvo aqui para info . Entre outras coisas, esse objeto de metadados inclui o número de exemplos de treinamento e teste.

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

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

Definir estratégia de distribuição

Crie um objeto MirroredStrategy . Isso tratará da distribuição e fornecerá um gerenciador de contexto ( tf.distribute.MirroredStrategy.scope ) para construir seu modelo interno.

strategy = tf.distribute.MirroredStrategy()

WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
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

Configurar pipeline de entrada

Ao treinar um modelo com várias GPUs, você pode usar o poder de computação extra de forma eficaz, aumentando o tamanho do lote. Em geral, use o maior tamanho de lote que cabe na memória da GPU e ajuste a taxa de aprendizado de acordo.

# 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

Os valores de pixel, que são 0-255, devem ser normalizados para o intervalo 0-1 . Defina esta escala em uma função.

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

  return image, label

Aplique esta função aos dados de treinamento e teste, embaralhe os dados de treinamento e agrupe-os para treinamento . Observe que também mantemos um cache na memória dos dados de treinamento para melhorar o desempenho.

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

Crie o modelo

Crie e compile o modelo Keras no contexto 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'])

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

Defina os retornos de chamada

Os retornos de chamada usados aqui são:

TensorBoard : este retorno de chamada grava um registro para o TensorBoard que permite visualizar os gráficos.
Ponto de verificação do modelo : este retorno de chamada salva o modelo após cada época.
Programador de taxa de aprendizado : usando este retorno de chamada, você pode programar a taxa de aprendizado para mudar após cada época / lote.

Para fins ilustrativos, adicione um retorno de chamada de impressão para exibir a taxa de aprendizagem no bloco de notas.

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

Treine e avalie

Agora, treine o modelo da maneira usual, chamando o fit no modelo e passando o conjunto de dados criado no início do tutorial. Esta etapa é a mesma, quer você esteja distribuindo o treinamento ou não.

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',).
3/938 [..............................] - ETA: 3:57 - loss: 2.2014 - accuracy: 0.2292WARNING:tensorflow:Callback method `on_train_batch_begin` is slow compared to the batch time (batch time: 0.0043s vs `on_train_batch_begin` time: 0.0693s). Check your callbacks.
WARNING:tensorflow:Callback method `on_train_batch_begin` is slow compared to the batch time (batch time: 0.0043s vs `on_train_batch_begin` time: 0.0693s). Check your callbacks.
WARNING:tensorflow:Callback method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0043s vs `on_train_batch_end` time: 0.0141s). Check your callbacks.
WARNING:tensorflow:Callback method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0043s vs `on_train_batch_end` time: 0.0141s). Check your callbacks.
938/938 [==============================] - 8s 4ms/step - loss: 0.1970 - accuracy: 0.9419

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

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

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

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

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

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

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

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

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

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

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

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

Como você pode ver abaixo, os pontos de verificação estão sendo salvos.

# 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

Para ver o desempenho do modelo, carregue o último ponto de verificação e chame evaluate nos dados de teste.

Chame a evaluate como antes, usando conjuntos de dados apropriados.

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 [==============================] - 2s 3ms/step - loss: 0.0383 - accuracy: 0.9870
Eval loss: 0.0383150540292263, Eval Accuracy: 0.9869999885559082

Para ver o resultado, você pode fazer o download e visualizar os registros do TensorBoard no terminal.

$ tensorboard --logdir=path/to/log-directory

ls -sh ./logs

total 4.0K
4.0K train

Exportar para SavedModel

Exporte o gráfico e as variáveis para o formato SavedModel independente de plataforma. Depois que seu modelo for salvo, você pode carregá-lo com ou sem o osciloscópio.

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

Carregue o modelo sem 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 [==============================] - 0s 2ms/step - loss: 0.0383 - accuracy: 0.9870
Eval loss: 0.0383150540292263, Eval Accuracy: 0.9869999885559082

Carregue o modelo com 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 2ms/step - loss: 0.0383 - accuracy: 0.9870
Eval loss: 0.0383150540292263, Eval Accuracy: 0.9869999885559082

Exemplos e tutoriais

Aqui estão alguns exemplos de uso de estratégia de distribuição com keras fit / compile:

Exemplo de transformador treinado usando tf.distribute.MirroredStrategy
Exemplo de NCF treinado usando tf.distribute.MirroredStrategy .

Mais exemplos listados no guia de estratégia de distribuição

Próximos passos

Leia o guia de estratégia de distribuição .
Leia o tutorial Treinamento distribuído com loops de treinamento personalizados .
Visite a seção Desempenho do guia para saber mais sobre outras estratégias e ferramentas que você pode usar para otimizar o desempenho de seus modelos do TensorFlow.