TF_CONFIG e treinamento distribuído

Quando você executa um job de treinamento, o AI Platform Training define uma variável de ambiente denominada TF_CONFIG em cada instância de máquina virtual (VM, na sigla em inglês) que faz parte do job. Seu código de treinamento, que é executado em cada VM, pode usar a variável de ambiente TF_CONFIG para acessar detalhes sobre o job de treinamento e o papel da VM em que está sendo executada.

O TensorFlow usa a variável de ambiente TF_CONFIG para facilitar o treinamento distribuído, mas você provavelmente não precisa acessá-lo diretamente em seu código de treinamento. Neste documento, descrevemos a variável de ambiente TF_CONFIG e respectivo uso em jobs distribuídos do TensorFlow e jobs de ajuste de hiperparâmetros.

Formato de TF_CONFIG

O AI Platform Training define a variável de ambiente TF_CONFIG em cada VM de cada job de treinamento para atender às especificações exigidas pelo TensorFlow para treinamento distribuído (em inglês). No entanto, o AI Platform Training também define outros campos na variável de ambiente TF_CONFIG além dos exigidos pelo TensorFlow.

A variável de ambiente TF_CONFIG é uma string JSON com o seguinte formato:

Campos de TF_CONFIG
cluster

Descrição do cluster do TensorFlow. Um dicionário que correlaciona um ou mais nomes de tarefa (chief, worker, ps ou master) com listas de endereços de rede em que essas tarefas estão em execução. Para um determinado job de treinamento, esse dicionário é o mesmo em todas as VMs.

Este é um primeiro argumento válido para o construtor tf.train.ClusterSpec. Observe que este dicionário nunca contém evaluator como chave, já que os avaliadores não são considerados parte do cluster de treinamento, mesmo se você usá-los para o job.

Saiba mais sobre a diferença entre chief e master em outra seção deste documento.
task

A descrição da tarefa da VM em que esta variável de ambiente está definida. Para um determinado job de treinamento, esse dicionário é diferente em cada VM. É possível usar essas informações para personalizar o código executado em cada VM em um job de treinamento distribuído. Também é possível usá-las para alterar o comportamento de seu código de treinamento para diferentes testes de um job de ajuste de hiperparâmetros.

Esse dicionário inclui os seguintes pares de chave-valor:

Campos task
type

O tipo de tarefa que a VM está executando. Esse valor é definido como worker nos workers, ps nos servidores de parâmetros e evaluator nos avaliadores. No worker mestre do job, o valor é definido como chief ou master. Saiba mais sobre a diferença entre os dois na seção chief x master deste documento.
index

O índice de tarefa com base em zero. Por exemplo, se seu job de treinamento incluir dois workers, este valor será definido como 0 em um deles e 1 no outro.
trial

O ID do teste de ajuste de hiperparâmetros atualmente em execução nesta VM. Este campo só será definido se o job de treinamento atual for um job de ajuste de hiperparâmetros.

Para jobs de ajuste de hiperparâmetros, o AI Platform Training executa seu código de treinamento várias vezes em vários testes com hiperparâmetros diferentes a cada vez. Este campo contém o número do teste atual, começando com 1 para o primeiro teste.
cloud

Um ID usado internamente pelo AI Platform Training. É possível ignorar este campo.
job

O TrainingInput que você forneceu para criar o job de treinamento atual, representado como um dicionário.

environment

A string cloud.

Para jobs de treinamento de contêiner personalizados, o AI Platform Training define uma variável de ambiente extra denominada CLUSTER_SPEC, que tem formato semelhante a TF_CONFIG, mas com várias diferenças importantes. Saiba mais sobre a variável de ambiente CLUSTER_SPEC.

Exemplo

No código de exemplo a seguir, a variável de ambiente TF_CONFIG é impressa nos seus registros de treinamento:

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

Em um job de ajuste de hiperparâmetros executado no ambiente de execução versão 2.1 ou posterior e que usa um worker mestre, dois workers e um servidor de parâmetros, este código produz o seguinte registro para um dos workers durante o primeiro teste de ajuste de hiperparâmetros. O exemplo de saída oculta o campo job para concisão e substitui alguns IDs 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 x master

A VM do worker mestre no AI Platform Training corresponde ao tipo de tarefa chief no TensorFlow. O TensorFlow pode designar uma tarefa worker para agir como chief, o AI Platform Training sempre designa explicitamente um chief.

master é um tipo de tarefa obsoleto no TensorFlow. master representava uma tarefa com papel semelhante a chief, mas também funcionava como evaluator em algumas configurações. O TensorFlow 2 não é compatível com variáveis de ambiente TF_CONFIG que contenham uma tarefa master.

O AI Platform Training usa chief nos campos cluster e task da variável de ambiente TF_CONFIG quando alguma das seguintes condições é verdadeira:

Caso contrário, por motivos de compatibilidade, o AI Platform Training usa o tipo de tarefa master obsoleto em vez de chief.

Quando usar TF_CONFIG

Como mencionado em uma seção anterior, você provavelmente não precisa interagir com a variável de ambiente TF_CONFIG diretamente em seu código de treinamento. Acesse a variável de ambiente TF_CONFIG apenas se as estratégias de distribuição do TensorFlow e o fluxo de trabalho de ajuste de hiperparâmetros padrão do AI Platform Training, ambos descritos nas próximas seções, não funcionarem para seu job.

Treinamento distribuído

O AI Platform Training define a variável de ambiente TF_CONFIG para ampliar as especificações necessárias ao TensorFlow para treinamento distribuído (em inglês).

Para realizar treinamento distribuído com o TensorFlow, use a API tf.distribute.Strategy (em inglês). Em particular, recomendamos que você use a API Keras com o MultiWorkerMirroredStrategy ou, se você especificar servidores de parâmetros para seu job, o ParameterServerStrategy. No entanto, o TensorFlow atualmente só fornece compatibilidade experimental para essas estratégias.

Essas estratégias de distribuição usam a variável de ambiente TF_CONFIG para atribuir papéis a cada VM em seu job de treinamento e para facilitar a comunicação entre as VMs. Você não precisa acessar a variável de ambiente TF_CONFIG diretamente no seu código de treinamento, porque o TensorFlow cuida disso para você.

Somente analise a variável de ambiente TF_CONFIG diretamente se quiser personalizar o comportamento das diferentes VMs que executam seu job de treinamento.

Ajuste de hiperparâmetros

Quando você executa um job de ajuste de hiperparâmetros, O AI Platform Training fornece argumentos diferentes para o código de treinamento de cada teste. Seu código de treinamento não precisa necessariamente estar ciente de qual teste está sendo realizado. Além disso, o AI Platform Training fornece ferramentas para monitorar o progresso de jobs de ajuste de hiperparâmetros.

Se necessário, seu código poderá ler o número do teste atual no campo trial do campo task da variável de ambiente TF_CONFIG.

A seguir