TF_CONFIG et l'entraînement distribué

Lorsque vous exécutez une tâche d'entraînement, AI Platform Training définit une variable d'environnement appelée TF_CONFIG sur chaque instance de machine virtuelle (VM) qui fait partie de votre tâche. Votre code d'entraînement, qui s'exécute sur chaque VM, peut utiliser la variable d'environnement TF_CONFIG pour accéder aux détails de la tâche d'entraînement et au rôle de la VM sur laquelle il s'exécute.

TensorFlow utilise la variable d'environnement TF_CONFIG pour faciliter l'entraînement distribué, mais vous n'avez probablement pas besoin d'y accéder directement dans votre code d'entraînement. Ce document décrit la variable d'environnement TF_CONFIG et son utilisation dans les tâches TensorFlow distribuées et les tâches de réglage d'hyperparamètres.

Format de TF_CONFIG

AI Platform Training définit la variable d'environnement TF_CONFIG sur chaque VM de chaque tâche d'entraînement pour répondre aux spécifications requises par TensorFlow pour l'entraînement distribué. Toutefois, AI Platform Training définit également des champs supplémentaires dans la variable d'environnement TF_CONFIG, en plus de ceux requis par TensorFlow.

La variable d'environnement TF_CONFIG est une chaîne JSON au format suivant :

Champs TF_CONFIG
cluster

Description du cluster TensorFlow. Dictionnaire associant un ou plusieurs noms de tâches (chief, worker, ps ou master) à des listes d'adresses réseau où ces tâches sont exécutées. Pour un job d'entraînement donnée, ce dictionnaire est identique sur toutes les VM.

Il s'agit d'un premier argument valide pour le constructeur tf.train.ClusterSpec. Notez que ce dictionnaire ne contient jamais l'élément evaluator en tant que clé, car les évaluateurs ne sont pas considérés comme faisant partie du cluster d'entraînement, même si vous les utilisez pour votre tâche.

Découvrez la différence entre chief et master dans une autre section de ce document.
task

Description de la tâche de la VM sur laquelle cette variable d'environnement est définie. Pour un job d'entraînement donné, ce dictionnaire est identique sur toutes les VM. Vous pouvez vous servir de ces informations pour personnaliser le code exécuté sur chaque VM dans un job d'entraînement distribué. Vous pouvez également l'utiliser pour modifier le comportement de votre code d'entraînement pour différents essais d'un job de réglage d'hyperparamètres.

Ce dictionnaire inclut les paires clé-valeur suivantes :

Champs task
type

Type de tâche effectuée par cette VM. Cette valeur est définie sur worker pour les nœuds de calcul, sur ps pour les serveurs de paramètres et sur evaluator pour les évaluateurs. Pour le nœud de calcul maître de votre job, la valeur est définie sur chief ou master. Découvrez la différence entre chief et master dans une autre section de ce document.
index

Index basé sur zéro de la tâche. Par exemple, si votre job d'entraînement inclut deux nœuds de calcul, cette valeur est définie sur 0 sur l'un d'entre eux et sur 1 sur l'autre.
trial

ID de l'essai de réglage d'hyperparamètres en cours d'exécution sur cette VM. Ce champ n'est défini que si la tâche d'entraînement en cours est une tâche de réglage d'hyperparamètres.

Pour les tâches de réglage d'hyperparamètres, AI Platform Training exécute votre code d'entraînement de manière répétée dans de nombreux essais avec des hyperparamètres différents à chaque fois. Ce champ contient le numéro d'essai en cours et commence à 1 pour le premier essai.
cloud

ID utilisé en interne par AI Platform Training. Vous pouvez ne pas tenir compte de ce champ.
job

Le TrainingInput que vous avez fourni pour créer la tâche d'entraînement actuelle, représenté sous la forme d'un dictionnaire.

environment

Correspond à la chaîne cloud.

Pour les tâches d'entraînement de conteneurs personnalisés, AI Platform Training définit une variable d'environnement supplémentaire appelée CLUSTER_SPEC, dont le format est semblable à TF_CONFIG, mais avec quelques différences importantes. En savoir plus sur la variable d'environnement CLUSTER_SPEC.

Exemple

L'exemple de code suivant imprime la variable d'environnement TF_CONFIG dans vos journaux d'entraînement :

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

Dans un job de réglage d'hyperparamètres exécuté dans la version d'exécution 2.1 ou ultérieure et utilisant un nœud de calcul maître, deux nœuds de calcul et un serveur de paramètres, ce code génère le journal suivant pour l'un des nœuds de calcul lors du premier essai de réglage d'hyperparamètres. L'exemple de résultat masque le champ job pour plus de concision et remplace certains ID par des valeurs génériques.

{
  "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 et master

La VM du nœud de calcul maître dans AI Platform Training correspond au type de tâche chief dans TensorFlow. Alors que TensorFlow peut désigner une tâche worker pour agir en tant que chief, AI Platform Training désigne toujours explicitement un chief.

master est un type de tâche obsolète dans TensorFlow. master représentait une tâche qui jouait le même rôle que chief, mais qui agissait également en tant que evaluator dans certaines configurations. TensorFlow 2 n'est pas compatible avec les variables d'environnement TF_CONFIG contenant une tâche master.

AI Platform Training utilise chief dans les champs cluster et task de la variable d'environnement TF_CONFIG si l'une des conditions suivantes est remplie :

Sinon, pour des raisons de compatibilité, AI Platform Training utilise le type de tâche master obsolète au lieu de chief.

Dans quel contexte utiliser TF_CONFIG ?

Comme indiqué dans une section précédente, vous n'avez probablement pas besoin d'interagir directement avec la variable d'environnement TF_CONFIG dans votre code d'entraînement. N'accédez à la variable d'environnement TF_CONFIG que si les stratégies de distribution de TensorFlow et le workflow standard de réglage d'hyperparamètres d'AI Platform Training, décrits dans les sections suivantes, ne fonctionnent pas.

Entraînement distribué

AI Platform Training définit la variable d'environnement TF_CONFIG pour étendre les spécifications requises par TensorFlow pour l'entraînement distribué.

Pour effectuer un entraînement distribué avec TensorFlow, utilisez l'API tf.distribute.Strategy. En particulier, nous vous recommandons d'utiliser l'API Keras avec MultiWorkerMirroredStrategy ou, si vous spécifiez des serveurs de paramètres pour votre tâche, avec ParameterServerStrategy. Toutefois, notez que TensorFlow ne fournit actuellement qu'une compatibilité expérimentale pour ces stratégies.

Ces stratégies de distribution utilisent la variable d'environnement TF_CONFIG pour attribuer des rôles à chaque VM dans votre job d'entraînement et pour faciliter la communication entre les VM. Vous n'avez pas besoin d'accéder à la variable d'environnement TF_CONFIG directement dans votre code d'entraînement, car TensorFlow s'en charge pour vous.

N'analysez la variable d'environnement TF_CONFIG directement que si vous souhaitez personnaliser le comportement des différentes VM exécutant votre job d'entraînement.

Réglages d'hyperparamètres

Lorsque vous exécutez une tâche de réglage d'hyperparamètres, AI Platform Training fournit différents arguments à votre code d'entraînement pour chaque essai. Votre code d'entraînement ne doit pas nécessairement savoir quel essai est en cours. En outre, AI Platform Training fournit des outils permettant de surveiller la progression des tâches de réglage des hyperparamètres.

Si nécessaire, votre code peut lire le numéro d'essai en cours dans le champ trial du champ task de la variable d'environnement TF_CONFIG.

Étapes suivantes