TF_CONFIG y entrenamiento distribuido

Cuando ejecutas un trabajo de entrenamiento, AI Platform Training configura una variable de entorno llamada TF_CONFIG en cada instancia de máquina virtual (VM) que forma parte de tu trabajo. El código de entrenamiento, que se ejecuta en cada VM, puede usar la variable de entorno TF_CONFIG para acceder a los detalles sobre el trabajo de entrenamiento y la función de la VM en la que se ejecuta.

TensorFlow usa la variable de entorno TF_CONFIG para facilitar el entrenamiento distribuido, pero es probable que no tengas que acceder a este directamente en el código de entrenamiento. En este documento, se describe la variable de entorno TF_CONFIG y su uso en trabajos distribuidos de TensorFlow y en trabajos de ajuste de hiperparámetros.

El formato de TF_CONFIG

AI Platform Training establece la variable de entorno TF_CONFIG en cada VM de cada trabajo de entrenamiento a fin de cumplir con las especificaciones que TensorFlow requiere para el entrenamiento distribuido. Sin embargo, AI Platform Training también establece más campos en la variable de entorno TF_CONFIG de los que requiere TensorFlow.

La variable de entorno TF_CONFIG es una string JSON con el siguiente formato:

Campos TF_CONFIG
cluster	La descripción del clúster de TensorFlow. Un diccionario que asigna uno o más nombres de tareas (chief, worker, ps o master) a las listas de direcciones de red en las que se ejecutan estas tareas. Para un trabajo de entrenamiento determinado, este diccionario es el mismo en todas las VM. Este es un primer argumento válido para el constructor tf.train.ClusterSpec. Ten en cuenta que este diccionario nunca contiene evaluator como clave, ya que los evaluadores no se consideran parte del clúster de entrenamiento, incluso si los usas para el trabajo. Obtén más información sobre la diferencia entre chief y master en otra sección de este documento.
task	La descripción de la tarea de la VM en la que se configura esta variable de entorno. Para un trabajo de entrenamiento determinado, este diccionario es diferente en cada VM. Puedes usar esta información para personalizar el código que se ejecuta en cada VM en un trabajo de entrenamiento distribuido. También puedes usarla para cambiar el comportamiento del código de entrenamiento en diferentes pruebas de un trabajo de ajuste de hiperparámetros. En este diccionario, se incluyen los siguientes pares clave-valor: Campos task type El tipo de tarea que realiza esta VM. Este valor se establece en worker en los trabajadores, ps en los servidores de parámetros y evaluator en los evaluadores. En el trabajador principal del trabajo, el valor se establece en chief o master. Obtén más información sobre la diferencia entre los dos en la sección chief en comparación con master de este documento. index El índice basado en cero de la tarea. Por ejemplo, si el trabajo de entrenamiento incluye dos trabajadores, este valor se establece en 0 en uno de ellos y en 1 en el otro. trial El ID de la prueba de ajuste de hiperparámetros que se está ejecutando en el momento en esta VM. Este campo solo se configura si el trabajo de entrenamiento actual es un trabajo de ajuste de hiperparámetros. En el caso de los trabajos de ajuste de hiperparámetros, AI Platform Training ejecuta el código de entrenamiento varias veces en varias pruebas con hiperparámetros diferentes cada vez. Este campo contiene el número de prueba actual, que comienza en 1 para la primera prueba. cloud Un ID que usa AI Platform Training de forma interna. Puedes ignorar este campo.
job	La TrainingInput que proporcionaste para crear el trabajo de entrenamiento actual, representada como un diccionario.
environment	La string cloud.

En los trabajos de entrenamiento de contenedores personalizados, AI Platform Training establece una variable de entorno adicional llamada CLUSTER_SPEC, que tiene un formato similar a TF_CONFIG pero con varias diferencias importantes. Obtén más información sobre la variable de entorno CLUSTER_SPEC.

Ejemplo

Con el siguiente código de ejemplo, se imprime la variable de entorno TF_CONFIG en los registros de entrenamiento:

import json
import os

tf_config_str = os.environ.get('TF_CONFIG')
tf_config_dict  = json.loads(tf_config_str)

# Convert back to string just for pretty printing
print(json.dumps(tf_config_dict, indent=2))

En un trabajo de ajuste de hiperparámetros que se ejecuta en la versión 2.1 o posterior del entorno de ejecución y usa un trabajador principal, dos trabajadores y un servidor de parámetros, este código produce el siguiente registro para uno de los trabajadores durante la primera prueba de ajuste de hiperparámetros. El resultado de ejemplo oculta el campo job para mayor concisión y reemplaza algunos ID por valores genéricos.

{
  "cluster": {
    "chief": [
      "cmle-training-chief-[ID_STRING_1]-0:2222"
    ],
    "ps": [
      "cmle-training-ps-[ID_STRING_1]-0:2222"
    ],
    "worker": [
      "cmle-training-worker-[ID_STRING_1]-0:2222",
      "cmle-training-worker-[ID_STRING_1]-1:2222"
    ]
  },
  "environment": "cloud",
  "job": {
    ...
  },
  "task": {
    "cloud": "[ID_STRING_2]",
    "index": 0,
    "trial": "1",
    "type": "worker"
  }
}

chief en comparación con master

La VM de trabajador principal en AI Platform Training corresponde al tipo de tarea chief en TensorFlow. Si bien TensorFlow puede designar una tarea worker para que actúe como chief, AI Platform Training siempre designa chief de manera explícita.

master es un tipo de tarea obsoleta en TensorFlow. master representaba una tarea que realizaba una función similar a chief, pero también actuaba como evaluator en algunas opciones de configuración. TensorFlow 2 no es compatible con las variables de entorno TF_CONFIG que contienen una tarea master.

AI Platform Training usa chief en los campos cluster y task de la variable de entorno TF_CONFIG si se cumple alguna de las siguientes condiciones:

• Ejecutas un trabajo de entrenamiento que usa versión 2.1 o posterior del entorno de ejecución.
• Configuraste el trabajo de entrenamiento para que use uno o más evaluadores. En otras palabras, configuraste el trainingInput.evaluatorCount del trabajo en 1 o superior.
• Tu trabajo usa un contenedor personalizado, y estableciste la trainingInput.useChiefInTfConfig del trabajo en true.

De lo contrario, por razones de compatibilidad, AI Platform Training usa el tipo de tarea master obsoleta en lugar de chief.

Cuándo usar TF_CONFIG

Como se mencionó en la sección anterior, es probable que no necesites interactuar con la variable de entorno TF_CONFIG directamente en el código de entrenamiento. Solo accede a la variable de entorno TF_CONFIG si las estrategias de distribución de TensorFlow y el flujo de trabajo de ajuste de hiperparámetros estándar de AI Platform Training, descritos en las siguientes secciones, no funcionan en tu trabajo.

Entrenamiento distribuido

AI Platform Training establece la variable de entorno TF_CONFIG a fin de ampliar las especificaciones que TensorFlow requiere para el entrenamiento distribuido.

Para realizar un entrenamiento distribuido con TensorFlow, usa la API tf.distribute.Strategy. En particular, te recomendamos que uses la API de Keras junto con la MultiWorkerMirroredStrategy o, si especificas los servidores de parámetros para el trabajo, la ParameterServerStrategy. Sin embargo, ten en cuenta que, en la actualidad, TensorFlow solo proporciona asistencia experimental para estas estrategias.

Estas estrategias de distribución usan la variable de entorno TF_CONFIG para asignar funciones a cada VM en el trabajo de entrenamiento y facilitar la comunicación entre las VM. No necesitas acceder a la variable de entorno TF_CONFIG directamente en el código de entrenamiento, ya que TensorFlow la controla por ti.

Solo analiza directamente la variable de entorno TF_CONFIG si deseas personalizar el comportamiento de las diferentes VM que ejecutan el trabajo de entrenamiento.

Ajuste de hiperparámetros

Cuando ejecutas un trabajo de ajuste de hiperparámetros, AI Platform Training proporciona argumentos diferentes al código de entrenamiento de cada prueba. No es necesario que el código de entrenamiento tenga en cuenta la prueba que se encuentra en ejecución. Además, AI Platform Training proporciona herramientas para supervisar el progreso de los trabajos de ajuste de hiperparámetros.

Si es necesario, el código puede leer el número de prueba actual del campo trial del campo task de la variable de entorno TF_CONFIG.

¿Qué sigue?

• Consulta el instructivo en la documentación de TensorFlow sobre el entrenamiento de varios trabajadores con Keras
• Obtén información sobre el entrenamiento distribuido con contenedores personalizados en AI Platform Training.
• Obtén información sobre cómo implementar el ajuste de hiperparámetros en los trabajos de entrenamiento.