distributed_training.ipynb

Kernel: Python 3

Copyright 2018 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 TensorFlow

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver fuente en GitHub

Descargar bloc de notas

Descripción general

tf.distribute.Strategy es una API de TensorFlow para distribuir el entrenamiento entre varias GPU, varias máquinas o TPU. Usando esta API, puede distribuir sus modelos existentes y el código de entrenamiento con cambios mínimos en el código.

tf.distribute.Strategy se ha diseñado teniendo en cuenta estas metas clave:

• Que sea fácil de usar y compatible con múltiples segmentos de usuarios, incluidos investigadores, ingenieros de aprendizaje automático, etc.

• Ofrecer un buen rendimiento desde el primer momento.

• De fácil cambio entre estrategias.

Puede distribuir el entrenamiento usando tf.distribute.Strategy con una API de alto nivel como Keras Model.fit, así como bucles de entrenamiento personalizados (y, en general, cualquier cálculo que utilice TensorFlow).

En TensorFlow 2.x, usted puede ejecutar sus programas de forma eager, o en un grafo usando tf.function. tf.distribute.Strategy intenta ser compatible con estos dos modos de ejecución, pero funciona mejor con tf.function. El modo eager sólo se recomienda con fines de depuración y no se admite para tf.distribute.TPUStrategy. Aunque esta guía se centra en el entrenamiento, esta API también puede usarse para distribuir la evaluación y la predicción en diferentes plataformas.

Usted puede usar tf.distribute.Strategy con muy pocos cambios en su código, porque los componentes subyacentes de TensorFlow han sido modificados para ser conscientes de la estrategia. Esto incluye variables, capas, modelos, optimizadores, métricas, resúmenes y puntos de verificación.

En esta guía, aprenderá sobre varios tipos de estrategias y cómo puede usarlas en diferentes situaciones. Para aprender a depurar problemas de rendimiento, consulte la guía Optimizar el rendimiento de la GPU de TensorFlow.

Nota: Para una comprensión más profunda de los conceptos, vea la presentación detallada Dentro de TensorFlow: tf.distribute.Strategy. Se recomienda especialmente si planea escribir su propio bucle de entrenamiento.

Preparación de TensorFlow

import tensorflow as tf

Tipos de estrategias

tf.distribute.Strategy intenta cubrir una serie de casos de uso a lo largo de diferentes ejes. Algunas de estas combinaciones ya están soportadas y otras se añadirán en el futuro. Algunos de estos ejes son:

• Entrenamiento síncrono vs asíncrono: Estas son dos formas comunes de distribuir el entrenamiento con paralelismo de datos. En el entrenamiento sincronizado, todos los trabajadores se capacitan sobre diferentes porciones de datos de entrada en sincronía, y agregando gradientes en cada paso. En el entrenamiento asíncrono, todos los trabajadores se entrenan independientemente sobre los datos de entrada y actualizan las variables de forma asíncrona. Normalmente, el entrenamiento sincronizado se admite a través de all-reduce y el asíncrono a través de la arquitectura de servidor de parámetros.

• Plataforma de hardware: Es posible que desee escalar su entrenamiento en varias GPU en una sola máquina, o varias máquinas en una red (con 0 o más GPU cada una), o en TPU en la nube.

Para poder usar estos casos, TensorFlow tiene MirroredStrategy, TPUStrategy, MultiWorkerMirroredStrategy, ParameterServerStrategy, CentralStorageStrategy, así como otras estrategias disponibles. La siguiente sección explica cuáles de ellas son soportadas en qué escenarios en TensorFlow. He aquí un resumen rápido:

API de entrenamiento	MirroredStrategy	TPUStrategy	MultiWorkerMirroredStrategy	CentralStorageStrategy	ParameterServerStrategy
Model.fit de Keras	Compatible	Compatible	Compatible	Compatible experimentalmente	Compatible experimentalmente
Bucle de entrenamiento personalizado	Compatible	Compatible	Compatible	Compatible experimentalmente	Compatible experimentalmente
API de estimador	Compatibilidad limitada	No compatible	Compatibilidad limitada	Compatibilidad limitada	Compatibilidad limitada

Nota: Compatible experimentalmente significa que las API no están cubiertas por ninguna garantía de compatibilidad.

Advertencia: El soporte del Estimator es limitado. El entrenamiento básico y la evaluación son experimentales, y las características avanzadas -como el andamiaje- no están implementadas. Debería usar Keras o bucles de entrenamiento personalizados si un caso de uso no está cubierto. No se recomienda el uso de Estimators para código nuevo. Los Estimators ejecutan código del estilo v1.Session, que es más difícil de escribir correctamente y puede comportarse de forma inesperada, especialmente cuando se combina con código TF 2. Los Estimators entran dentro de nuestras garantías de compatibilidad, pero no recibirán más parches que las vulnerabilidades de seguridad. Vaya a la guía de migración para más detalles.

MirroredStrategy

tf.distribute.MirroredStrategy admite el entrenamiento distribuido síncrono en varias GPU de una misma máquina. Crea una réplica por cada dispositivo GPU. Cada variable del modelo se refleja en todas las réplicas. Juntas, estas variables forman una única variable conceptual denominada MirroredVariable. Estas variables se conservan sincronizadas entre sí aplicando actualizaciones idénticas.

Se usan algoritmos eficientes tipo All-reduce para comunicar las actualizaciones de las variables entre los dispositivos. El All-reduce agrega los tensores en todos los dispositivos sumándolos, y los pone a disposición en cada dispositivo. Se trata de un algoritmo fusionado que es muy eficiente y puede reducir la sobrecarga de la sincronización de forma significativa. Hay muchos algoritmos e implementaciones de all-reduce disponibles, dependiendo del tipo de comunicación disponible entre dispositivos. De forma predeterminada, usa la librería de comunicación colectiva de NVIDIA (NCCL) como implementación de all-reduce. Puede seleccionar entre otras opciones o escribir la suya propia.

He aquí la forma más sencilla de crear MirroredStrategy:

mirrored_strategy = tf.distribute.MirroredStrategy()

Esto creará una instancia MirroredStrategy, que usará todas las GPU visibles para TensorFlow, y NCCL-como la comunicación entre dispositivos.

Si desea usar sólo algunas de las GPU de su máquina, puede hacerlo así:

mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])

Si desea omitir la comunicación entre dispositivos, puede hacerlo usando el argumento cross_device_ops indicando una instancia de tf.distribute.CrossDeviceOps. Actualmente, tf.distribute.HierarchicalCopyAllReduce y tf.distribute.ReductionToOneDevice son dos opciones distintas de tf.distribute.NcclAllReduce, que es la predeterminada.

mirrored_strategy = tf.distribute.MirroredStrategy(
    cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())

TPUStrategy

tf.distribute.TPUStrategy le permite ejecutar su entrenamiento de TensorFlow en Unidades de Procesamiento de Tensores (TPUs). Las TPUs son ASICs especializados de Google diseñadas para acelerar drásticamente las cargas de trabajo de aprendizaje automático. Están disponibles en Google Colab, la Nube de Investigación TPU, y Cloud TPU.

En términos de arquitectura de entrenamiento distribuido, TPUStrategy es la misma MirroredStrategy: implementa el entrenamiento distribuido síncrono. Las TPU ofrecen su propia implementación de all-reduce eficiente y otras operaciones colectivas a través de múltiples núcleos de TPU, que se usan en TPUStrategy.

Así es como se instanciaría TPUStrategy:

Nota: Para ejecutar cualquier código TPU en Colab, debe seleccionar TPU como el runtime de Colab. Consulte la guía Utilizar TPU para ver un ejemplo completo.

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
    tpu=tpu_address)
tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
tpu_strategy = tf.distribute.TPUStrategy(cluster_resolver)

La instancia TPUClusterResolver ayuda a localizar las TPU. En Colab, no es necesario especificarle ningún argumento.

Si quiere usar esto para Cloud TPUs:

• Debe especificar el nombre de su recurso TPU en el argumento tpu.

• Debe inicializar el sistema TPU explícitamente en el inicio del programa. Esto se requiere antes de poder usar las TPU para calcular. Inicializar el sistema TPU también borra la memoria TPU, por lo que es importante completar este paso primero para evitar la pérdida de estado.

MultiWorkerMirroredStrategy

tf.distribute.MultiWorkerMirroredStrategy es muy similar a MirroredStrategy. Implementa el entrenamiento distribuido síncrono a través de múltiples trabajadores, cada uno con potencialmente múltiples GPUs. Similar a tf.distribute.MirroredStrategy, crea copias de todas las variables del modelo en cada dispositivo a través de todos los trabajadores.

Esta es la forma más sencilla de crear MultiWorkerMirroredStrategy:

strategy = tf.distribute.MultiWorkerMirroredStrategy()

MultiWorkerMirroredStrategy tiene dos implementaciones para las comunicaciones entre dispositivos. CommunicationImplementation.RING está basada en RPC y admite tanto CPU como GPU. CommunicationImplementation.NCCL usa NCCL y proporciona un rendimiento de última generación en GPUs, pero no es compatible con CPUs. CollectiveCommunication.AUTO aplaza la elección a Tensorflow. Puede especificarlos de la siguiente manera:

communication_options = tf.distribute.experimental.CommunicationOptions(
    implementation=tf.distribute.experimental.CommunicationImplementation.NCCL)
strategy = tf.distribute.MultiWorkerMirroredStrategy(
    communication_options=communication_options)

Una de las diferencias clave para poner en marcha el entrenamiento multitrabajador, en comparación con el entrenamiento multiGPU, es la configuración multitrabajador. La variable de entorno 'TF_CONFIG' es la forma estándar en TensorFlow de especificar la configuración del cluster a cada trabajador que forma parte del cluster. Aprenda más en la sección 'Ajuste de TF_CONFIG' de este documento.

Si quiere conocer más detalles sobre MultiWorkerMirroredStrategy, consulte los siguientes tutoriales:

• Entrenamiento multitrabajador con modelo Keras.fit

• Entrenamiento multitrabajador con un bucle de entrenamiento personalizado

ParameterServerStrategy

El entrenamiento de servidores de parámetros es un método común de datos paralelos para escalar el entrenamiento de modelos en múltiples máquinas. Un clúster de entrenamiento con servidores de parámetros está formado por trabajadores y servidores de parámetros. Las variables se crean en los servidores de parámetros y son leídas y actualizadas por los trabajadores en cada paso. Consulte el tutorial Entrenamiento de servidores de parámetros para más detalles.

En TensorFlow 2, el entrenamiento del servidor de parámetros usa una arquitectura basada en un coordinador central a través de la clase tf.distribute.experimental.coordinator.ClusterCoordinator.

En esta implementación, las tareas del trabajador y del servidor de parámetros ejecutan tf.distribute.Servers que escuchan las tareas del coordinador. El coordinador crea recursos, envía tareas de entrenamiento, escribe puntos de verificación y se encarga de los fallos de las tareas.

En la programación que se ejecuta en el coordinador, utilizará un objeto ParameterServerStrategy para definir un paso de entrenamiento y usará un ClusterCoordinator para enviar los pasos de entrenamiento a los trabajadores remotos. Esta es la forma más sencilla de crearlos:

strategy = tf.distribute.experimental.ParameterServerStrategy(
    tf.distribute.cluster_resolver.TFConfigClusterResolver(),
    variable_partitioner=variable_partitioner)
coordinator = tf.distribute.experimental.coordinator.ClusterCoordinator(
    strategy)

Para saber más sobre ParameterServerStrategy, consulte el Tutorial de entrenamiento del servidor de parámetros con Keras Model.fit y un bucle de entrenamiento personalizado.

Nota: Deberá configurar la variable de entorno 'TF_CONFIG' si usa TFConfigClusterResolver. Es similar a 'TF_CONFIG' en MultiWorkerMirroredStrategy pero tiene salvedades adicionales.

En TensorFlow 1, ParameterServerStrategy sólo está disponible con un Estimator a través del símbolo tf.compat.v1.distribute.experimental.ParameterServerStrategy.

Nota: Esta estrategia es experimental ya que actualmente se encuentra en desarrollo activo.

CentralStorageStrategy

tf.distribute.experimental.CentralStorageStrategy también realiza el entrenamiento síncrono. Las variables no se replican, sino que se colocan en la CPU y las operaciones se replican en todas las GPU locales. Si sólo hay una GPU, todas las variables y operaciones se colocarán en esa GPU.

Cree una instancia de CentralStorageStrategy mediante:

central_storage_strategy = tf.distribute.experimental.CentralStorageStrategy()

Esto creará una instancia de CentralStorageStrategy que usará todas las GPU y CPU visibles. La actualización de las variables en las réplicas se agregará antes de aplicarse a las variables.

Nota: Esta estrategia es experimental, ya que actualmente se encuentra en fase de desarrollo.

Otras estrategias

Aparte de las estrategias anteriores, hay otras dos que pueden ser útiles para la creación de prototipos y la depuración cuando se usan APIs tf.distribute.

Estrategia predeterminada

La estrategia Default es una estrategia de distribución que se aplica cuando no existe ninguna estrategia de distribución explícita. Implementa la interfaz tf.distribute.Strategy pero es una transferencia y no proporciona ninguna distribución real. Por ejemplo, Strategy.run(fn) simplemente llamará a fn. El código escrito usando esta estrategia debería comportarse exactamente igual que el código escrito sin ninguna estrategia. Puede pensar en ella como una estrategia "a falta de decisión".

La estrategia Default es un único ejemplar y no se pueden crear más instancias de ella. Puede obtenerse utilizando tf.distribute.get_strategy fuera del ámbito de cualquier estrategia explícita (la misma API que puede usarse para obtener la estrategia actual dentro del ámbito de una estrategia explícita).

default_strategy = tf.distribute.get_strategy()

Esta estrategia tiene dos objetivos principales:

• Permite escribir código de librería consciente de la distribución de forma incondicional. Por ejemplo, en tf.keras.optimizers puede usar tf.distribute.get_strategy y usar esa estrategia para reducir gradientes (siempre devolverá un objeto de estrategia sobre el que puede llamar a la API Strategy.reduce).

# In optimizer or other library code
# Get currently active strategy
strategy = tf.distribute.get_strategy()
strategy.reduce("SUM", 1., axis=None)  # reduce some values

• Similar al código de librería, puede usarse para escribir programas de usuarios finales para trabajar con y sin estrategia de distribución, sin requerir lógica condicional. Este es un fragmento de código de muestra que lo ilustra:

if tf.config.list_physical_devices('GPU'):
  strategy = tf.distribute.MirroredStrategy()
else:  # Use the Default Strategy
  strategy = tf.distribute.get_strategy()

with strategy.scope():
  # Do something interesting
  print(tf.Variable(1.))

OneDeviceStrategy

tf.distribute.OneDeviceStrategy es una estrategia para colocar todas las variables y el cálculo en un único dispositivo especificado.

strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")

Esta estrategia se diferencia de la estrategia Default en varios aspectos. En la estrategia Default, la lógica de colocación de variables permanece sin cambios en comparación con la ejecución de TensorFlow sin ninguna estrategia de distribución. Pero cuando se usa OneDeviceStrategy, todas las variables creadas en su ámbito se colocan explícitamente en el dispositivo especificado. Además, cualquier función llamada a través de OneDeviceStrategy.run también se colocará en el dispositivo especificado.

La entrada distribuida a través de esta estrategia se preextraerá al dispositivo especificado. En la estrategia Default, no hay distribución de entrada.

De forma similar a la estrategia Default, esta estrategia también podría usarse para probar su código antes de cambiar a otras estrategias que realmente realicen la distribución a múltiples dispositivos/máquinas. Esto ejercitará la maquinaria de la estrategia de distribución algo más que la estrategia Default, pero no hasta el punto que lo haría usar, por ejemplo, MirroredStrategy o TPUStrategy. Si desea un código que se comporte como si no hubiera ninguna estrategia, use la estrategia Default.

Hasta ahora ha aprendido sobre las diferentes estrategias y cómo puede instanciarlas. Las siguientes secciones muestran las distintas formas en que puede usarlas para distribuir su entrenamiento.

Usar tf.distribute.Strategy con Keras Model.fit

tf.distribute.Strategy está integrada en tf.keras, que es la implementación de TensorFlow de la especificación Keras API. tf.keras es una API de alto nivel para construir y entrenar modelos. Al integrarse en el backend tf.keras, le resultará muy sencillo distribuir su entrenamiento escrito en el marco de entrenamiento Keras utilizando Model.fit.

Esto es lo que debe cambiar en su código:

1. Cree una instancia de la tf.distribute.Strategy adecuada.

2. Traslade la creación del modelo Keras, el optimizador y las métricas dentro de strategy.scope. De este modo, el código de los métodos call(), train_step() y test_step() del modelo se distribuirá y ejecutará en el acelerador o aceleradores.

Las estrategias de distribución de TensorFlow admiten todos los tipos de modelos Keras: secuencial, funcional y subclasificado.

Este es un fragmento de código para hacerlo para un modelo Keras muy simple con una capa Dense:

mirrored_strategy = tf.distribute.MirroredStrategy()

with mirrored_strategy.scope():
  model = tf.keras.Sequential([
      tf.keras.layers.Dense(1, input_shape=(1,),
                            kernel_regularizer=tf.keras.regularizers.L2(1e-4))])
  model.compile(loss='mse', optimizer='sgd')

Este ejemplo usa MirroredStrategy, por lo que puede ejecutarlo en una máquina con varias GPU. strategy.scope() indica a Keras qué estrategia usar para distribuir el entrenamiento. La creación de modelos/optimizadores/métricas dentro de este ámbito le permite crear variables distribuidas en lugar de variables normales. Una vez establecido esto, puede ajustar su modelo como lo haría normalmente. MirroredStrategy se encarga de replicar el entrenamiento del modelo en las GPU disponibles, agregar gradientes, etc.

dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(10)
model.fit(dataset, epochs=2)
model.evaluate(dataset)

Aquí un tf.data.Dataset facilita la entrada para el entrenamiento y la evaluación. También puede usar arreglos NumPy:

import numpy as np

inputs, targets = np.ones((100, 1)), np.ones((100, 1))
model.fit(inputs, targets, epochs=2, batch_size=10)

En ambos casos (con Dataset o NumPy), cada lote de la entrada dada se divide a partes iguales entre las múltiples réplicas. Por ejemplo, si está usando MirroredStrategy con 2 GPUs, cada lote de tamaño 10 se dividirá entre las 2 GPUs, recibiendo cada una 5 ejemplos de entrada en cada paso. Cada época se entrenará más rápido a medida que se añadan más GPU. Normalmente, querrá aumentar el tamaño de los lotes a medida que añada más aceleradores, para usar de forma eficaz la potencia de cálculo adicional. También tendrá que reajustar su tasa de aprendizaje, en función del modelo. Puede usar strategy.num_replicas_in_sync para obtener el número de réplicas.

mirrored_strategy.num_replicas_in_sync

# Compute a global batch size using a number of replicas.
BATCH_SIZE_PER_REPLICA = 5
global_batch_size = (BATCH_SIZE_PER_REPLICA *
                     mirrored_strategy.num_replicas_in_sync)
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100)
dataset = dataset.batch(global_batch_size)

LEARNING_RATES_BY_BATCH_SIZE = {5: 0.1, 10: 0.15, 20:0.175}
learning_rate = LEARNING_RATES_BY_BATCH_SIZE[global_batch_size]

Qué es compatible ahora

API de entrenamiento	MirroredStrategy	TPUStrategy	MultiWorkerMirroredStrategy	ParameterServerStrategy	CentralStorageStrategy
Model.fit de Keras	Compatible	Compatible	Compatible	Compatible experimentalmente	Compatible experimentalmente

Ejemplos y tutoriales

Aquí hay una lista de tutoriales y ejemplos que ilustran la integración anterior de extremo a extremo con Model.fit de Keras:

1. Tutorial: Entrenamiento con Model.fit y MirroredStrategy.

2. Tutorial: Entrenamiento con Model.fit y MultiWorkerMirroredStrategy.

3. Guía: Contiene un ejemplo de uso de Model.fit y TPUStrategy.

4. Tutorial: Entrenamiento del servidor de parámetros con Model.fit y ParameterServerStrategy.

5. Tutorial: Ajuste fino de BERT para muchas tareas del índice de referencia GLUE con Model.fit y TPUStrategy.

6. Repositorio Model Garden de TensorFlow, que contiene colecciones de modelos de última generación implementados usando diversas estrategias.

Usar tf.distribute.Strategy con bucles de entrenamiento personalizados

Como se ha demostrado anteriormente, usar tf.distribute.Strategy con Model.fit de Keras requiere cambiar sólo un par de líneas de su código. Con un poco más de esfuerzo, también puede usar tf.distribute.Strategy con bucles de entrenamiento personalizados.

Si necesita más flexibilidad y control sobre sus bucles de entrenamiento de lo que es posible con Estimator o Keras, puede escribir bucles de entrenamiento personalizados. Por ejemplo, al usar un GAN, puede que desee dar un número diferente de pasos de generador o discriminador en cada ronda. Igualmente, los marcos de alto nivel no son muy adecuados para el entrenamiento del aprendizaje por refuerzo.

Las clases tf.distribute.Strategy ofrecen un conjunto básico de métodos para soportar bucles de entrenamiento personalizados. Usarlos puede requerir una reestructuración menor del código inicialmente, pero una vez hecho esto, debería ser capaz de cambiar entre GPUs, TPUs y múltiples máquinas simplemente cambiando la instancia de la estrategia.

Abajo encontrará un breve fragmento que ilustra este caso de uso para un sencillo ejemplo de entrenamiento usando el mismo modelo Keras que antes.

En primer lugar, cree el modelo y el optimizador dentro del ámbito de la estrategia. Esto asegura que cualquier variable creada con el modelo y el optimizador sean variables espejo.

with mirrored_strategy.scope():
  model = tf.keras.Sequential([
      tf.keras.layers.Dense(1, input_shape=(1,),
                            kernel_regularizer=tf.keras.regularizers.L2(1e-4))])
  optimizer = tf.keras.optimizers.SGD()

A continuación, cree el conjunto de datos de entrada y llame a tf.distribute.Strategy.experimental_distribute_dataset para distribuir el conjunto de datos en función de la estrategia.

dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(1000).batch(
    global_batch_size)
dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)

Luego, defina un paso del entrenamiento. Use tf.GradientTape para calcular gradientes y el optimizador para aplicar esos gradientes para actualizar las variables de su modelo. Para distribuir este paso de entrenamiento, póngalo en una función train_step y pásela a tf.distribute.Strategy.run junto con las entradas del conjunto de datos que obtuvo del dist_dataset creado anteriormente:

# Sets `reduction=NONE` to leave it to tf.nn.compute_average_loss() below.
loss_object = tf.keras.losses.BinaryCrossentropy(
  from_logits=True,
  reduction=tf.keras.losses.Reduction.NONE)

def train_step(inputs):
  features, labels = inputs

  with tf.GradientTape() as tape:
    predictions = model(features, training=True)
    per_example_loss = loss_object(labels, predictions)
    loss = tf.nn.compute_average_loss(per_example_loss)
    model_losses = model.losses
    if model_losses:
      loss += tf.nn.scale_regularization_loss(tf.add_n(model_losses))

  gradients = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))
  return loss

@tf.function
def distributed_train_step(dist_inputs):
  per_replica_losses = mirrored_strategy.run(train_step, args=(dist_inputs,))
  return mirrored_strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses,
                         axis=None)

Algunas otras cosas a tener en cuenta en el código anterior:

1. Ha usado tf.nn.compute_average_loss para reducir las pérdidas de predicción por cada ejemplo a un escalar. tf.nn.compute_average_loss suma la pérdida por ejemplo y divide la suma por el tamaño global del lote. Esto es importante porque más tarde, después de hacer el cálculo de los gradientes en cada réplica, son agregados a través de las réplicas sumándolos.

De forma predeterminada, el tamaño del lote global se toma como tf.get_strategy().num_replicas_in_sync * tf.shape(per_example_loss)[0]. También puede especificarse explícitamente como argumento de palabra clave global_batch_size=. Sin lotes cortos, lo predeterminado equivale a tf.nn.compute_average_loss(..., global_batch_size=global_batch_size) con el global_batch_size definido anteriormente. (Para obtener más información sobre los lotes cortos y cómo evitarlos o gestionarlos, consulte el Tutorial de entrenamiento personalizado.)

1. Ha usado tf.nn.scale_regularization_loss para escalar las pérdidas de regularización registradas con el objeto Model, si las hubiera, también por 1/num_replicas_in_sync. Para las pérdidas de regularización que dependen de la entrada, corresponde al código de modelado, y no al bucle de entrenamiento personalizado, realizar el promedio sobre el tamaño del lote por réplica(!); de este modo, el código de modelado puede permanecer agnóstico de la réplica mientras que el bucle de entrenamiento permanece agnóstico de cómo se calculan las pérdidas de regularización.

2. Cuando se llama a apply_gradients dentro del ámbito de una estrategia de distribución, se modifica su comportamiento. En concreto, antes de aplicar gradientes en cada instancia paralela durante el entrenamiento síncrono, realiza una suma sobre todas las réplicas de los gradientes.

3. También ha usado la API tf.distribute.Strategy.reduce para agregar los resultados devueltos por tf.distribute.Strategy.run para generar informes. tf.distribute.Strategy.run devuelve resultados de cada réplica local de la estrategia, y existen múltiples formas de consumir este resultado. Puede usar reduce (reducirlos) para obtener un valor agregado. También puede utilizar tf.distribute.Strategy.experimental_local_results para obtener la lista de valores contenidos en el resultado, uno por cada réplica local.

Por último, una vez definido el paso de entrenamiento, puede iterar sobre dist_dataset y ejecutar el entrenamiento en un bucle:

for dist_inputs in dist_dataset:
  print(distributed_train_step(dist_inputs))

En el ejemplo anterior, usted iteró sobre el dist_dataset para aportar entradas a su entrenamiento. También se le facilita el tf.distribute.Strategy.make_experimental_numpy_dataset para dar soporte a las entradas NumPy. Puede usar esta API para crear un conjunto de datos antes de llamar a tf.distribute.Strategy.experimental_distribute_dataset.

Otra forma de iterar sobre sus datos es usar explícitamente iteradores. Es posible que desee hacer esto cuando desee ejecutar un número determinado de pasos en lugar de iterar sobre todo el conjunto de datos. La iteración anterior se modificaría ahora para crear primero un iterador y luego llamar explícitamente next sobre él para obtener los datos de entrada.

iterator = iter(dist_dataset)
for _ in range(10):
  print(distributed_train_step(next(iterator)))

Esto cubre el caso más sencillo de usar la API tf.distribute.Strategy para distribuir bucles de entrenamiento personalizados.

Qué es compatible ahora

API de entrenamiento	MirroredStrategy	TPUStrategy	MultiWorkerMirroredStrategy	ParameterServerStrategy	CentralStorageStrategy
Bucle de entrenamiento personalizado	Compatible	Compatible	Compatible	Compatible experimentalmente	Compatible experimentalmente

Ejemplos y tutoriales

Aquí tiene algunos ejemplos para usar estrategias de distribución con bucles de entrenamiento personalizados:

1. Tutorial: Entrenamiento con un bucle de entrenamiento personalizado y MirroredStrategy.

2. Tutorial: Entrenamiento con un bucle de entrenamiento personalizado y MultiWorkerMirroredStrategy.

3. Guía: Contiene un ejemplo de bucle de entrenamiento personalizado con TPUStrategy.

4. Tutorial: Entrenamiento del servidor de parámetros con un bucle de entrenamiento personalizado y ParameterServerStrategy.

5. Repositorio Model Garden de TensorFlow, que contiene colecciones de modelos de última generación implementados usando diversas estrategias.

Otros temas

Esta sección cubre algunos temas que son relevantes para múltiples casos de uso.

Ajuste de la variable de entorno TF_CONFIG

Para el entrenamiento multitrabajador, como ya se ha mencionado, necesita ajustar la variable de entorno 'TF_CONFIG' para cada binario que se ejecute en su cluster. La variable de entorno 'TF_CONFIG' es una cadena JSON que especifica qué tareas constituyen un clúster, sus direcciones y el rol de cada tarea en el clúster. La repo tensorflow/ecosystem provee una plantilla Kubernetes, que configura 'TF_CONFIG' para sus tareas de entrenamiento.

Hay dos componentes de 'TF_CONFIG': un cluster y una tarea.

• Un cluster ofrece información sobre el cluster de entrenamiento, que es un dict formado por diferentes tipos de trabajos como trabajadores. En el entrenamiento multitrabajador, suele haber un trabajador que asume algo más de responsabilidad, como guardar el punto de verificación y escribir el archivo de resumen para TensorBoard, además de lo que hace un trabajador normal. A dicho trabajador se le refiere como el trabajador "jefe", y es habitual que el trabajador con índice 0 sea designado como trabajador jefe (de hecho, así es como se implementa tf.distribute.Strategy).

• Por otro lado, una tarea ofrece información sobre la tarea actual. El primer componente cluster es el mismo para todos los trabajadores, y el segundo componente task es diferente en cada trabajador y especifica el tipo y el índice de ese trabajador.

Un ejemplo de 'TF_CONFIG' es:

os.environ["TF_CONFIG"] = json.dumps({
    "cluster": {
        "worker": ["host1:port", "host2:port", "host3:port"],
        "ps": ["host4:port", "host5:port"]
    },
   "task": {"type": "worker", "index": 1}
})

Este 'TF_CONFIG' especifica que hay tres trabajadores y dos tareas "ps" en el "cluster" junto con sus hosts y puertos. La parte "task" especifica el rol de la tarea actual en el "cluster": trabajador 1 (el segundo trabajador). Los roles válidos en un cluster son "chief", "worker", "ps", y "evaluator". No debería haber ningún trabajo "ps" excepto cuando se usa tf.distribute.experimental.ParameterServerStrategy.

¿Qué sigue?

tf.distribute.Strategy está activamente en desarrollo. Pruébelo y aporte su retroalimentación usando GitHub issues.