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.

Entrenamiento distribuido con Keras

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver fuente en GitHub

Descargar el bloc de notas

Descripción general

La API tf.distribute.Strategy ofrece una abstracción para distribuir su entrenamiento entre varias unidades de procesamiento. Permite realizar un entrenamiento distribuido usando los modelos y el código de entrenamiento existentes con cambios mínimos.

Este tutorial muestra cómo usar tf.distribute.MirroredStrategy para implementar la replicación dentro del grafo con entrenamiento síncrono en muchas GPU de una misma máquina. Esencialmente, la estrategia copia todas las variables del modelo en cada procesador. Luego usa all-reduce para combinar los gradientes de todos los procesadores y aplica el valor combinado a todas las copias del modelo.

Usará las API de tf.keras API para construir el modelo y Model.fit para entrenarlo (si quiere aprender sobre la capacitación distribuida con un bucle de entrenamiento personalizado y el MirroredStrategy, consulte este tutorial).

MirroredStrategy entrene su modelo en varias GPU de una sola máquina. Para entrenamiento síncrono en muchas GPUs en múltiples trabajadores, usa el tf.distribute.MultiWorkerMirroredStrategy con el Keras Model.fit o un bucle de capacitación personalizado. Para otras opciones, refiérete a la Guía de entrenamiento distribuido.

Para conocer otras estrategias, consulte la guía Entrenamiento distribuido con TensorFlow.

Preparación

import tensorflow_datasets as tfds
import tensorflow as tf

import os

# Load the TensorBoard notebook extension.
%load_ext tensorboard

print(tf.__version__)

Descargar el conjunto de datos

Cargue el conjunto de datos MNIST de Conjuntos de datos TensorFlow. Esto devuelve un conjunto de datos en formato tf.data.

Si establece el argumento con_info en True, se incluirán los metadatos de todo el conjunto de datos, que se almacenan aquí en info. Este objeto de metadatos incluye, entre otras cosas, el número de ejemplos de entrenamiento y de prueba.

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

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

Definir la estrategia de distribución

Crea un objeto MirroredStrategy. Se encargará de la distribución y creará un gestor de contexto (MirroredStrategy.scope) en el que construir tu modelo.

strategy = tf.distribute.MirroredStrategy()

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

Configurar la canalización de entrada

Si está entrenando un modelo con varias GPU, puedes usar eficazmente la potencia de procesamiento adicional aumentando el tamaño del lote. Por lo general, debes usar el tamaño de lote más grande que quepa en la memoria de la GPU y ajustar la tasa de aprendizaje en consecuencia.

# 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

Define una función que normalice los valores de los pixeles de imagen que van de [0, 255] a [0, 1] (escalado de características):

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

  return image, label

Aplica esta función scale a los datos de entrenamiento y de prueba, y luego usa las APIs tf.data.Dataset para mezclar los datos de entrenamiento (Dataset.shuffle) y procesarlos por lotes (Dataset.batch). Recuerda conservar también una caché en memoria de los datos de entrenamiento para mejorar el rendimiento (Dataset.cache).

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

Crear el modelo e instanciar el optimizador

En el contexto de Strategy.scope, crea y compila el modelo usando la API 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 ejemplo ficticio con el conjunto de datos MNIST, usará la tasa de aprendizaje predeterminada del optimizador Adam de 0.001.

Para conjuntos de datos más grandes, el beneficio clave del entrenamiento distribuido es aprender más en cada paso del entrenamiento, porque cada paso procesa más datos de entrenamiento en paralelo, lo que permite una mayor tasa de aprendizaje (dentro de los límites del modelo y del conjunto de datos).

Definir las retrollamadas

Define los siguientes retrollamadas Keras:

• tf.keras.callbacks.TensorBoard: escribe un log para TensorBoard, que te permite visualizar los grafos.

• tf.keras.callbacks.ModelCheckpoint: guarda el modelo con cierta frecuencia, por ejemplo, después de cada época.

• tf.keras.callbacks.BackupAndRestore: añade la funcionalidad de tolerancia a fallos haciendo una copia de seguridad del modelo y del número de época actual. Para más información, consulta la sección Tolerancia a fallos del tutorial Entrenamiento multitrabajador con Keras.

• tf.keras.callbacks.LearningRateScheduler: planifica el cambio del ritmo de aprendizaje después de, por ejemplo, cada época/lote.

Para ilustrarlo, añade una retrollamada personalizada denominada PrintLR para mostrar la velocidad de aprendizaje en el bloc de notas.

Nota: Usa la retrollamada BackupAndRestore en lugar de ModelCheckpoint como mecanismo principal para restaurar el estado del entrenamiento cuando falle un trabajo. Como BackupAndRestore sólo funciona en eager mode, considera usar ModelCheckpoint en modo de 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()
]

Entrenar y evaluar

Ahora, entrena el modelo del modo habitual llamando a Model.fit de Keras en el modelo y pasándole el conjunto de datos creado al principio del tutorial. Este paso es el mismo tanto si distribuyes el entrenamiento como si no.

EPOCHS = 12

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

Busca puntos de verificación guardados:

# Check the checkpoint directory.
!ls {checkpoint_dir}

Carga el último punto de verificación y llama a Model.evaluate con los datos de prueba para ver el rendimiento del modelo:

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 el resultado, inicia TensorBoard y mira los logs:

%tensorboard --logdir=logs

!ls -sh ./logs

Guardar el modelo

Guarda el modelo en un archivo zip .keras usando Model.save. Cuando tengas el modelo guardado, puedes cargarlo con o sin el Strategy.scope.

path = 'my_model.keras'

model.save(path)

Ahora carga el modelo sin 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))

Carga el modelo con 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 adicionales

Más ejemplos que usan diferentes estrategias de distribución con la API Model.fit de Keras:

1. El tutorial Resolver tareas GLUE utilizando BERT en TPU usa tf.distribute.MirroredStrategy para entrenar en GPUs y tf.distribute.TPUStrategy en TPUs.

2. El tutorial Guardar y cargar un modelo mediante una estrategia de distribución muestra cómo usar las API SavedModel con tf.distribute.Strategy.

3. Los modelos oficiales de TensorFlow pueden configurarse para ejecutar múltiples estrategias de distribución.

Para saber más sobre las estrategias de distribución de TensorFlow:

1. El tutorial Entrenamiento personalizado con tf.distribute.Strategy enseña a usar tf.distribute.MirroredStrategy para entrenar a un solo trabajador con un bucle de entrenamiento personalizado.

2. El tutorial Entrenamiento multitrabajador con Keras enseña a usar la MultiWorkerMirroredStrategy con Model.fit.

3. El tutorial Bucle de entrenamiento personalizado con Keras y MultiWorkerMirroredStrategy muestra cómo usar la MultiWorkerMirroredStrategy con Keras y un bucle de entrenamiento personalizado.

4. La guía Entrenamiento distribuido en TensorFlow da una visión general de las estrategias de distribución disponibles.

5. La guía Mejor rendimiento con tf.function ofrece información sobre otras estrategias y herramientas, como el Perfilador TensorFlow que puedes usar para optimizar el rendimiento de tus modelos TensorFlow.

Nota: tf.distribute.Strategy se encuentra activamente en desarrollo y TensorFlow añadirá más ejemplos y tutoriales en un futuro próximo. Anímate a probarlo. Tu retroalimentación es bienvenida, no dudes en enviarla a través de issues en GitHub.