keras.ipynb

Kernel: Python 3

Copyright 2019 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Treinamento distribuído com o Keras

Ver em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

Visão geral

A API tf.distribute.Strategy fornece uma abstração para a distribuição do seu treinamento em diversas unidades de processamento. Ela permite realizar treinamento distribuído usando modelos e código de treinamento existentes, com mudanças mínimas.

Este tutorial demonstra como usar tf.distribute.MirroredStrategy para fazer uma replicação no grafo com treinamento síncrono em diversas GPUs de uma máquina. Basicamente, a estratégia copia todas as variáveis do modelo em cada processador. Em seguida, é feita uma redução total (all-reduce) para combinar os gradientes de todos os processadores, e o valor combinado é aplicado a todas as cópias do modelo.

Você usará as APIs tf.keras para construir o modelo e Model.fit para treiná-lo. (Para saber mais sobre treinamento distribuído com um loop de treinamento personalizado e o MirroredStrategy , confira este tutorial .)

MirroredStrategy treina seu modelo em várias GPUs em uma única máquina. Para treinamento síncrono em várias GPUs em vários funcionários , use tf.distribute.MultiWorkerMirroredStrategy com Keras Model.fit ou um loop de treinamento personalizado . Para outras opções, consulte o guia de treinamento distribuído .

Para saber mais sobre diversas outras estratégias, confira o guia Treinamento distribuído com o TensorFlow.

Configuração

import tensorflow_datasets as tfds
import tensorflow as tf

import os

# Load the TensorBoard notebook extension.
%load_ext tensorboard

print(tf.__version__)

Download do dataset

Baixe o dataset MNIST nos Datasets do TensorFlow. Você baixará um dataset no formato tf.data.

A definição do argumento with_info como True (verdadeiro) inclui os metadados de todo o dataset, que estão sendo salvos em info. Entre outras coisas, esse objeto de metadados inclui o número dos 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']

Definição da estratégia de distribuição

Crie um objeto MirroredStrategy. Isso fará a distribuição e fornecerá um gerenciador de contexto (MirroredStrategy.scope) para criar seu modelo.

strategy = tf.distribute.MirroredStrategy()

print('Number of devices: {}'.format(strategy.num_replicas_in_sync))

Configuração do pipeline de entrada

Ao treinar um modelo com diversas GPUs, você pode aumentar o tamanho do lote para usar o poder computacional extra. De forma geral, use o maior tamanho de lote que cabe na memória das GPUs e ajuste e taxa de aprendizado de acordo com ele.

# 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

Defina uma função que normalize os valores de pixel de imagens do intervalo [0, 255] para o intervalo [0, 1] (escalonamento de características):

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

  return image, label

Aplique essa função de scale aos dados de treinamento e teste. Depois, use as APIs tf.data.Dataset para misturar os dados de treinamento (Dataset.shuffle) e divida em lotes (Dataset.batch). Note que você também está mantendo um cache dos dados de treinamento na memória para aumentar o desempenho (Dataset.cache).

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 e instancie o otimizador

No contexto de strategy.scope(), crie e compile o modelo usando a API do Keras:

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(learning_rate=0.001),
                metrics=['accuracy'])

Para este exemplo de brinquedo com o dataset MNIST, você usará a taxa de aprendizado padrão do otimizador Adam, que é 0,001.

Para datasets maiores, o principal benefício do treinamento distribuído é aprender mais em cada passo do treinamento, porque cada passo processa mais dados de treinamento em paralelo, o que permite uma taxa de aprendizagem maior (dentro dos limites do modelo e do dataset).

Definição dos callbacks

Defina os seguintes callbacks do Keras:

• tf.keras.callbacks.TensorBoard: escreve um log para o TensorBoard, permitindo que você visualize os grafos.

• tf.keras.callbacks.ModelCheckpoint: salva o modelo com uma determinada frequência, como após cada época.

• tf.keras.callbacks.BackupAndRestore: faz backup do modelo e do número da época atual para fornecer a funcionalidade de tolerância a falhas. Saiba mais na seção Tolerância a falhas do tutorial Treinamento multiworker com o Keras.

• tf.keras.callbacks.LearningRateScheduler: agenda a alteração da taxa de aprendizado após cada época/lote, por exemplo

Para fins ilustrativos, adicione um callback personalizado chamado PrintLR para exibir a taxa de aprendizado no bloco de notas

Observação: use o callback BackupAndRestore em vez de ModelCheckpoint como o mecanismo principal para restaurar o estado de treinamento ao reiniciar após a falha de um trabalho. Como BackupAndRestore é compatível somente com o modo adiantado (eager), considere usar ModelCheckpoint no modo grafo.

# Define the checkpoint directory to store the checkpoints.
checkpoint_dir = './training_checkpoints'
# Define the name of the checkpoint files.
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt_{epoch}")

# Define a 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

# Define a callback for printing the learning rate 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()))

# Put all the callbacks together.
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()
]

Treinamento e avaliação

Agora, treine o modelo da maneira habitual: faça uma chamada ao Model.fit do Keras no modelo e passe o dataset criado no começo do tutorial. Este passo é sempre igual, não importa se você esteja distribuindo o treinamento ou não.

EPOCHS = 12

model.fit(train_dataset, epochs=EPOCHS, callbacks=callbacks)

Verifique os checkpoints salvos:

# Check the checkpoint directory.
!ls {checkpoint_dir}

Para verificar o desempenho do modelo, carregue o último checkpoint e faça uma chamada a Model.evaluate nos dados de teste:

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

Para visualizar a saída, inicie o TensorBoard e confira os logs:

%tensorboard --logdir=logs

!ls -sh ./logs

Salvando o modelo

Salve o modelo como um arquivo zip .keras usando Model.save. Após salvar o modelo, você pode carregá-lo com ou sem o strategy.scope().

path = 'my_model.keras'

model.save(path)

Agora, 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))

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

Recursos adicionais

Confira mais exemplos que usam diferentes estratégias de distribuição com a API Model.fit do Keras:

1. O tutorial Resolva tarefas GLUE usando BERT e TPU usa tf.distribute.MirroredStrategy para treinamento em GPUs e tf.distribute.TPUStrategy em TPUs.

2. O tutorial Salve e carregue um modelo usando uma estratégia de distribuição demonstra como usar as APIs SavedModel com tf.distribute.Strategy.

3. Os modelos oficiais do TensorFlow podem ser configurados para executarem diversas estratégias de distribuição.

Para saber mais sobre as estratégias de distribuição do TensorFlow:

1. O tutorial Treinamento personalizado com tf.distribute.Strategy mostra como usar tf.distribute.MirroredStrategy para treinamento em um único worker com um loop de treinamento personalizado.

2. O tutorial Treinamento multiworker com o Keras mostra como usar MultiWorkerMirroredStrategy com Model.fit.

3. O tutorial Loop de treinamento personalizado com o Keras e MultiWorkerMirroredStrategy mostra como usar MultiWorkerMirroredStrategy com o Keras e um loop de treinamento personalizado.

4. O guia Treinamento distribuído no TensorFlow apresenta uma visão geral das estratégias de distribuição disponíveis.

5. O guia Desempenho melhor com tf.function apresenta informações sobre outras estratégias e ferramentas, como o TensowFlower Profiler, que você pode usar para otimizar o desempenho dos seus modelos do TensorFlow.

Observação: tf.distribute.Strategy está em constante desenvolvimento, e o TensorFlow adicionará mais exemplos e tutoriais no futuro próximo. Faça um teste. Seu feedback será bem-vindo. Fique à vontade para enviá-lo pelo recurso Issues (Problemas) do GitHub.