Entrena en pods de TPU

Resumen

Uno de los beneficios de Cloud TPU es que puedes escalar el entrenamiento entre diferentes tipos de TPU. Todos los tipos de TPU usan la misma arquitectura de datos paralelos. El único cambio es que el paralelismo aumenta de 8 núcleos a 2,048 núcleos.

Para aprovechar al máximo las cantidades más grandes de TPU, debes ajustar varios parámetros de la tarea de entrenamiento. En este documento, se explican algunos problemas comunes y los cambios que debes realizar en tus modelos.

Escala el tamaño del lote y los pasos del entrenamiento

Para lograr el escalamiento lineal en los tipos de TPU más grandes, conserva el mismo tamaño del lote por núcleo.

Por ejemplo, si usas un tamaño del lote de 1,024 en una v2-8, usa un tamaño del lote de 4,096 (4 * 1,024) en un v2-32. Esto utiliza por completo el hardware de TPU. Puedes usar tamaños del lote más pequeños, pero tu entrenamiento no se escalará de forma lineal si lo haces.

Varios modelos incluyen una marca train_steps en la que un paso corresponde a un solo lote de datos. Cuando aumentes el tamaño del lote, reduce la cantidad de pasos de entrenamiento para que la cantidad total de ejemplos de entrenamiento permanezca igual.

Por ejemplo, si la configuración original era un tamaño de lote de 1,024 para 1,000 pasos, podrías aumentar el tamaño del lote a 4096 y reducir los pasos a 250 (1,000 / 4). Si tu modelo usa una marca epochs, no es necesario que escales la cantidad de pasos.

Los tamaños del lote más grandes pueden cambiar el comportamiento de convergencia del modelo, por lo que también puedes ajustar algunos hiperparámetros.

Usa optimizadores de lotes grandes

Si escalas verticalmente a fin de usar tamaños del lote más grandes, es posible que debas usar un optimizador que admita tamaños del lote grandes. Por ejemplo, el modelo de referencia de Resnet usa el LARSOptimizer para tipos de TPU más grandes.

Usa depósitos de Google Cloud Storage regionales en la misma región que el pod de TPU

En general, la práctica recomendada para el entrenamiento de la TPU es usar siempre recursos en la misma región. La región de recursos es muy importante cuando se usan pods de TPU, ya que la tasa de transferencia de datos de Google Cloud Storage tiende a ser más alta. Asegúrate de usar un depósito de Google Cloud Storage regional en la misma región que la TPU para los conjuntos de datos de entrenamiento y los puntos de control.

Proporciona el nombre y la región de TPU a TPUClusterResolver

Para el entrenamiento de un solo dispositivo, puedes especificar el nombre de TPU o una dirección IP, por ejemplo: grpc://1.2.3.4:8470.

Para los pods de TPU, debes usar el nombre de TPU a fin de que TensorFlow pueda descubrir las direcciones IP de todos los hosts disponibles para la distribución del entrenamiento.

Si especificas la dirección IP para un pod de TPU genera el error siguiente:

ValueError: TPUConfig.num_shards is not set correctly ....

Si necesitas ejecutar tu modelo y TPU en diferentes zonas o regiones, debes proporcionar TPUClusterResolver con la zona de la TPU.

Puedes pasar las marcas tpu y tpu_zone a los modelos de referencia o configurar directamente los valores para TPUClusterResolver. Por ejemplo:

cluster = tf.contrib.cluster_resolver.TPUClusterResolver(
  tpu='tpuname',
  zone='europe-west4-a',
  project='myproject')

Un problema relacionado que puede surgir cuando se usa TPUEstimator es no configurar el agente de resolución del clúster en tf.contrib.tpu.RunConfig. Esto podría suceder si RunConfig se inicializa con el parámetro master. En su lugar, debes inicializarlo con el parámetro cluster.

config = tf.contrib.tpu.RunConfig(cluster=cluster)

Si usas un pod de TPU y tf.Session sin usar TPUEstimator, debes establecer la configuración de la sesión de la siguiente manera:

with tf.Session(cluster.get_master(), config=config.session_config) as sess:

Evaluación en pods de TPU

La evaluación no es compatible ni rentable en los pods de TPU. Google recomienda ejecutar la evaluación en una TPU de un solo dispositivo con la marca mode, que es compatible con la mayoría de los modelos.

Para entrenar en un pod de TPU o una TPU de un solo dispositivo, configura la marca como train. Si deseas realizar la evaluación, configura eval solo para las TPU de un solo dispositivo. También puedes usar el modo train_and_eval en las TPU de un solo dispositivo, pero no en pods de TPU.