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.

Treinamento distribuído com o TensorFlow

Ver em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

Visão geral

tf.distribute.Strategy é uma API do TensorFlow para distribuir treinamento em múltiplas GPUs, múltiplas máquinas ou TPUs. Usando esta API, você pode distribuir seus modelos existentes e código de treinamento com alterações mínimas no código.

tf.distribute.Strategy foi projetada com os seguintes objetivos principais em mente:

• Ser fácil de usar e oferecer suporte a diversos segmentos de usuários, incluindo pesquisadores, engenheiros de aprendizado de máquina, etc.

• Fornecer bom desempenho pronto para uso.

• Ser fácil alternar entre estratégias.

Você pode distribuir o treinamento usando tf.distribute.Strategy com uma API de alto nível como Keras Model.fit, bem como loops de treinamento personalizados (e, em geral, qualquer computação usando TensorFlow).

No TensorFlow 2.x, você pode executar seus programas de forma eager ou num grafo usando tf.function. tf.distribute.Strategy pretende oferecer suporte a ambos os modos de execução, mas funciona melhor com tf.function. O modo eager é recomendado apenas para fins de depuração e não é compatível com tf.distribute.TPUStrategy. Embora o treinamento seja o foco deste guia, esta API também pode ser usada para distribuir avaliação e previsão em diferentes plataformas.

Você pode usar tf.distribute.Strategy com bem poucas alterações no seu código, porque os componentes do TensorFlow foram alterados para reconhecerem a estratégia usada. Isto inclui variáveis, camadas, modelos, otimizadores, métricas, resumos e checkpoints.

Neste guia, você aprenderá sobre vários tipos de estratégias e como usá-las em diferentes situações. Para saber como depurar problemas de desempenho, veja o guia Otimize o desempenho da GPU com o TensorFlow.

Observação: para uma compreensão mais profunda dos conceitos, assista à apresentação detalhada — Por dentro do TensorFlow: tf.distribute.Strategy. Isto é recomendado principalmente se você pretende escrever seu próprio loop de treinamento.

Configure o TensorFlow

import tensorflow as tf

Tipos de estratégias

tf.distribute.Strategy pretende cobrir uma série de casos de uso em diferentes eixos. Algumas dessas combinações são suportadas no momento e outras serão adicionadas no futuro. Alguns desses eixos são:

• Treinamento síncrono versus assíncrono: essas são duas formas comuns de distribuição de treinamento com paralelismo de dados. No treinamento síncrono, todos os workers treinam sobre diferentes fatias de dados de entrada de forma sincronizada e agregando gradientes em cada passo. No treinamento assíncrono, todos os workers treinam de forma independente sobre os dados de entrada e atualizam as variáveis ​​de forma assíncrona. Geralmente, o treinamento síncrono é suportado através de all-reduce e o treinamento assíncrono através da arquitetura do servidor de parâmetros.

• Plataforma de hardware: vocÊ talvez queira escalonar seu treinamento em múltiplas GPUs numa mesma máquina ou em diversas máquinas numa rede (com 0 ou mais GPUs cada) ou em TPUs na nuvem.

Para dar suporte a esses casos de uso, o TensorFlow disponibiliza MirroredStrategy, TPUStrategy, MultiWorkerMirroredStrategy, ParameterServerStrategy, CentralStorageStrategy, além de outras estratégias. A próxima seção explica quais delas são compatíveis com quais cenários no TensorFlow. Eis uma visão geral rápida:

API de treinamento	MirroredStrategy	TPUStrategy	MultiWorkerMirroredStrategy	CentralStorageStrategy	ParameterServerStrategy
Keras Model.fit	Suportado	Suportado	Suportado	Suporte experimental	Suporte experimental
Loop de treinamento personalizado	Suportado	Suportado	Suportado	Suporte experimental	Suporte experimental
Estimator API	Suporte limitado	Não suportado	Suporte limitado	Suporte limitado	Suporte limitado

Observação: suporte experimental significa que as APIs não são cobertas por nenhuma garantia de compatibilidade.

Importante: o suporte ao estimador é limitado. O treinamento básico e a avaliação são experimentais e recursos avançados – como o scaffold – não são implementados. Você deve usar Keras ou loops de treinamento personalizados se um caso de uso não for coberto. O uso de estimadores não é recomendado para novos códigos. Os estimadores executam código no estilo v1.Session que é mais difícil de escrever corretamente e pode se comportar de forma inesperada, especialmente quando combinado com código TF 2. Os estimadores se enquadram nas nossas garantias de compatibilidade, mas não receberão nenhuma correção exceto as relacionadas a vulnerabilidades de segurança. Veja o guia de migração para mais detalhes.

MirroredStrategy

tf.distribute.MirroredStrategy oferece suporte ao treinamento distribuído síncrono em múltiplas GPUs de uma máquina. Ela cria uma réplica por dispositivo GPU. Cada variável no modelo é espelhada em todas as réplicas. Juntas, essas variáveis ​​formam uma única variável conceitual chamada MirroredVariable. Essas variáveis ​​são mantidas sincronizadas entre si através da aplicação de atualizações idênticas.

Algoritmos all-reduce eficientes são usados ​​para comunicar as atualizações de variáveis ​​entre os dispositivos. All-reduce agrega tensores em todos os dispositivos, somando-os e disponibilizando-os em cada dispositivo. É um algoritmo de fusão que é muito eficiente e pode reduzir significativamente a sobrecarga da sincronização. Existem muitos algoritmos e implementações de all-reduce disponíveis, dependendo do tipo de comunicação disponível entre os dispositivos. Por padrão, ele usa a NVIDIA Collective Communication Library (NCCL) como implementação de all-reduce. Você pode escolher outra implementação ou escrever a sua própria.

Eis a maneira mais simples de criar uma MirroredStrategy:

mirrored_strategy = tf.distribute.MirroredStrategy()

Isso criará uma instância de MirroredStrategy, que usará todas as GPUs visíveis para o TensorFlow e o NCCL, para a comunicação entre dispositivos.

Se você pretender usar apenas algumas GPUs na sua máquina, você pode fazer da seguinte forma:

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

Se quiser substituir a comunicação entre dispositivos, você pode fazer isso usando o argumento cross_device_ops fornecendo uma instância de tf.distribute.CrossDeviceOps. Atualmente, tf.distribute.HierarchicalCopyAllReduce e tf.distribute.ReductionToOneDevice são duas opções diferentes de tf.distribute.NcclAllReduce, que é o padrão.

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

TPUStrategy

tf.distribute.TPUStrategy permite que você execute seu treinamento do TensorFlow em Unidades de Processamento de Tensor (TPUs). TPUs são ASICs especializadas do Google, projetadas para acelerar drasticamente as cargas de trabalho de aprendizado de máquina. Elas estão disponíveis no Google Colab, na TPU Research Cloud e na TPU em nuvem.

Em termos de arquitetura de treinamento distribuído, a TPUStrategy é a mesma MirroredStrategy: ela implementa treinamento distribuído síncrono. As TPUs fornecem sua própria implementação de operações eficientes de redução total e outras operações coletivas em mútiplos núcleos de TPU, que são usados ​​na TPUStrategy.

Eis como você instanciaria uma TPUStrategy:

Observação: para executar qualquer código TPU no Colab, você precisa selecionar TPU como o runtime do Colab. Consulte o guia Use TPUs para ver um exemplo 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.experimental.TPUStrategy(cluster_resolver)

A instância TPUClusterResolver ajuda a localizar as TPUs. No Colab, você não precisa especificar nenhum argumento.

Se você quiser usar isso para TPUs na nuvem:

• Você precisa passar o nome do seu recurso TPU como o argumento tpu.

• Você precisa inicializar o sistema TPU explicitamente no início do programa. Isto é necessário antes que as TPUs possam ser usadas para computação. A inicialização do sistema TPU também apaga a memória da TPU, por isso é importante concluir esta etapa primeiro para evitar a perda de estado.

MultiWorkerMirroredStrategy

tf.distribute.MultiWorkerMirroredStrategy é muito semelhante a MirroredStrategy. Ela implementa treinamento distribuído síncrono entre diversos workers, cada um com potencialmente várias GPUs. Semelhante a tf.distribute.MirroredStrategy, ele cria cópias de todas as variáveis ​​no modelo em cada dispositivo em todos os workers.

Aqui está a maneira mais simples de criar uma MultiWorkerMirroredStrategy:

strategy = tf.distribute.MultiWorkerMirroredStrategy()

MultiWorkerMirroredStrategy possui duas implementações para comunicações entre dispositivos. CommunicationImplementation.RING é baseado em RPC e oferece suporte a CPUs e GPUs. CommunicationImplementation.NCCL usa NCCL e fornece desempenho de última geração em GPUs, mas não oferece suporte a CPUs. CollectiveCommunication.AUTO adia a escolha para o Tensorflow. Você pode especificá-las da seguinte maneira:

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

Uma das principais diferenças em iniciar o treinamento multiworker, em comparação ao treinamento multi-GPU, é a configuração multiworker. A variável de ambiente 'TF_CONFIG' é a forma padrão no TensorFlow para especificar a configuração do cluster para cada worker que faz parte do cluster. Saiba mais na seção seção de configuração do TF_CONFIG deste documento.

Para obter mais detalhes sobre MultiWorkerMirroredStrategy, considere os seguintes tutoriais:

• Treinamento multiworker com Keras Model.fit

• Treinamento multiworker com um loop de treinamento personalizado

ParameterServerStrategy

O treinamento do servidor de parâmetros é um método de paralelismo de dados comum para ampliar o treinamento de modelos em várias máquinas. Um cluster de treinamento de servidores de parâmetros consiste em workers e servidores de parâmetros. As variáveis ​​são criadas em servidores de parâmetros e são lidas e atualizadas pelos workers em cada passo. Veja o tutorial Treinamento do servidor de parâmetros para mais detalhes.

No TensorFlow 2, o treinamento do servidor de parâmetros usa uma arquitetura baseada em coordenador central por meio da classe tf.distribute.experimental.coordinator.ClusterCoordinator.

Nesta implementação, as tarefas do worker e do parameter server executam tf.distribute.Server que escutam tarefas do coordenador. O coordenador cria recursos, despacha tarefas de treinamento, escreve checkpoints e lida com falhas de tarefas.

Na programação em execução no coordenador, você usará um objeto ParameterServerStrategy para definir um passo de treinamento e usará um ClusterCoordinator para despachar passos de treinamento para workers remotos. Aqui está a maneira mais simples de criá-los:

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

Para saber mais sobre a ParameterServerStrategy, veja o tutorial Treinamento de um servidor de parâmetros com Keras Model.fit e um loop de treinamento personalizado.

Observação: Você precisará configurar a variável de ambiente 'TF_CONFIG' se usar TFConfigClusterResolver. Ela é semelhante a 'TF_CONFIG' em MultiWorkerMirroredStrategy, mas tem ressalvas adicionais.

No TensorFlow 1, a ParameterServerStrategy está disponível apenas com um Estimador por meio do símbolo tf.compat.v1.distribute.experimental.ParameterServerStrategy.

Observação: Esta estratégia é experimental, pois está atualmente em desenvolvimento ativo.

CentralStorageStrategy

tf.distribute.experimental.CentralStorageStrategy também faz treinamento síncrono. As variáveis ​​não são espelhadas; em vez disso, são colocadas na CPU e as operações são replicadas em todas as GPUs locais. Se houver apenas uma GPU, todas as variáveis ​​e operações serão colocadas nessa GPU.

Crie uma instância do CentralStorageStrategy:

central_storage_strategy = tf.distribute.experimental.CentralStorageStrategy()

Isto criará uma instância CentralStorageStrategy que usará todas as GPUs e CPU visíveis. A atualização nas variáveis ​​nas réplicas será agregada antes de ser aplicada às variáveis.

Observação: Esta estratégia é experimental, pois atualmente é um trabalho em andamento.

Outras estratégias

Além das estratégias acima, existem duas outras estratégias que podem ser úteis para prototipagem e depuração ao usar APIs tf.distribute.

Default Strategy

A Default Strategy (estratégia padrão) é uma estratégia de distribuição que está presente quando nenhuma estratégia de distribuição explícita está no escopo. Ela implementa a interface tf.distribute.Strategy, mas passa direto, simplesmente retornando a entrada sem fornecer qualquer distribuição. Por exemplo, Strategy.run(fn) simplesmente chamará fn. O código escrito usando esta estratégia deve se comportar exatamente como o código escrito sem qualquer estratégia. Você pode pensar nela como uma estratégia “no-op”.

A Default Strategy é um singleton – e não é possível criar mais instâncias dela. Pode ser obtida usando tf.distribute.get_strategy fora do escopo de qualquer estratégia explícita (a mesma API que pode ser usada para colocar a estratégia atual dentro do escopo de uma estratégia explícita).

default_strategy = tf.distribute.get_strategy()

Esta estratégia tem duas finalidades principais:

• Ela permite escrever incondicionalmente código de API compatível com a distribuição. Por exemplo, em tf.keras.optimizers você pode usar tf.distribute.get_strategy e usar essa estratégia para reduzir gradientes: ela sempre retornará um objeto Strategy que você pode usar para chamar a 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

• De forma similar ao código de API, ela pode ser usada para escrever programas de usuários finais para trabalhar com ou sem uma estratégia de distribuição, sem exigir nenhuma lógica condicional. Aqui está um trecho de código de exemplo que ilustra isso:

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 é uma estratégia para reunir todas as variáveis ​​e computações num único dispositivo especificado.

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

Esta estratégia difere da Default Strategy de várias maneiras. Na Default Strategy, a lógica de posicionamento de variáveis ​​permanece inalterada quando comparada à execução do TensorFlow sem qualquer estratégia de distribuição. Mas ao usar OneDeviceStrategy, todas as variáveis ​​criadas no seu escopo são colocadas explicitamente no dispositivo especificado. Além disso, quaisquer funções chamadas via OneDeviceStrategy.run também serão colocadas no dispositivo especificado.

A entrada distribuída através desta estratégia será pré-buscada no dispositivo especificado. Na Default Strategy, não há distribuição da entrada.

Semelhante à Default Strategy, esta estratégia também pode ser usada para testar seu código antes de mudar para outras estratégias que realmente distribuem para múltiplos dispositivos/máquinas. Isto exercitará o mecanismo da estratégia de distribuição um pouco mais do que a Default Strategy, mas não ao máximo obtido pelo uso de, por exemplo, MirroredStrategy ou TPUStrategy. Se você deseja um código que se comporte como se não houvesse uma estratégia, use a Default Strategy.

Até aqui você aprendeu sobre diferentes estratégias e como instanciá-las. As próximas seções mostram as diferentes maneiras pelas quais você pode usá-las para distribuir seu treinamento.

Use tf.distribute.Strategy com Keras Model.fit

tf.distribute.Strategy está integrado ao tf.keras, que é a implementação da especificação da API Keras do TensorFlow. tf.keras é uma API de alto nível para construir e treinar modelos. Ao integrar-se ao back-end tf.keras, é fácil distribuir seu treinamento escrito no framework de treinamento Keras usando Model.fit.

Eis o que você precisa alterar em seu código:

1. Crie uma instância do tf.distribute.Strategy apropriado.

2. Mova a criação do modelo Keras, otimizador e métricas dentro de strategy.scope. Assim, o código nos métodos call(), train_step() e test_step() do modelo serão todos distribuídos e executados no(s) acelerador(es).

As estratégias de distribuição do TensorFlow oferecem suporte a todos os tipos de modelos Keras: sequencial, funcional e subclasse

Aqui está um trecho de código para fazer isso para um modelo Keras muito simples com uma camada 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 exemplo usa MirroredStrategy, então você pode executar o código numa máquina com múltiplas GPUs. strategy.scope() informa Keras sobre qual estratégia usar para distribuir o treinamento. A criação de modelos/otimizadores/métricas dentro deste escopo permite criar variáveis ​​distribuídas em vez de variáveis ​​comuns. Depois de configurado, você pode ajustar seu modelo como faria normalmente. MirroredStrategy se encarrega de replicar o treinamento do modelo nas GPUs disponíveis, agregando gradientes e muito mais.

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

Aqui, um tf.data.Dataset fornece a entrada de treinamento e avaliação. Você também pode usar arrays NumPy:

import numpy as np

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

Em ambos os casos – com Dataset ou NumPy – cada lote da entrada fornecida é dividido igualmente entre as múltiplas réplicas. Por exemplo, se você estiver usando MirroredStrategy com 2 GPUs, cada lote de tamanho 10 será dividido entre as 2 GPUs, com cada uma recebendo 5 exemplos de entrada em cada passo. Cada época será treinada mais rapidamente à medida que você adiciona mais GPUs. Normalmente, você vai querer aumentar o tamanho do lote à medida que adiciona mais aceleradores, para fazer uso eficaz do poder de computação adicional. Você também precisará reajustar sua taxa de aprendizado, dependendo do modelo. Você pode usar strategy.num_replicas_in_sync para obter o 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]

O que suportado agora?

API de treinamento	MirroredStrategy	TPUStrategy	MultiWorkerMirroredStrategy	ParameterServerStrategy	CentralStorageStrategy
Keras Model.fit	Suportado	Suportado	Suportado	Suporte experimental	Suporte experimental

Exemplos e tutoriais

Aqui está uma lista de tutoriais e exemplos que ilustram a integração acima de ponta a ponta com o Keras Model.fit:

1. Tutorial: Treinamento com Model.fit e MirroredStrategy.

2. Tutorial: Treinamento com Model.fit e MultiWorkerMirroredStrategy.

3. Guia: Contém um exemplo de uso de Model.fit e TPUStrategy.

4. Tutorial: Treinamento de servidor de parâmetros com Model.fit e ParameterServerStrategy.

5. Tutorial : Ajuste fino no BERT para muitas tarefas do benchmark GLUE com Model.fit e TPUStrategy.

6. Repositório TensorFlow Model Garden contendo coleções de modelos de última geração implementados usando diversas estratégias.

Use tf.distribute.Strategy com loops de treinamento personalizados

Conforme demonstrado acima, usar tf.distribute.Strategy com Keras Model.fit requer a alteração de apenas algumas linhas do seu código. Com um pouco mais de esforço, você também pode usar tf.distribute.Strategy com loops de treinamento personalizados.

Se você precisar de mais flexibilidade e controle sobre seus loops de treinamento do que seria possível com o Estimator ou Keras, você pode escrever loops de treinamento personalizados. Por exemplo, ao usar uma GAN, você pode querer realizar um número diferente de passos de gerador ou discriminador a cada rodada. Da mesma forma, os frameworks de alto nível não são muito adequados para o treinamento de Aprendizado por Reforço.

As classes tf.distribute.Strategy fornecem um conjunto básico de métodos para oferecer suporte a loops de treinamento personalizados. Usá-los pode exigir uma pequena reestruturação do código inicialmente, mas uma vez feito isso, você poderá alternar entre GPUs, TPUs e várias máquinas simplesmente alterando a instância da estratégia.

Abaixo está um pequeno trecho de código que ilustra este caso de uso para um exemplo de treinamento simples usando o mesmo modelo Keras de antes.

Primeiro, crie o modelo e o otimizador dentro do escopo da estratégia. Isso vai garantir que quaisquer variáveis ​​criadas com o modelo e o otimizador sejam variáveis ​​espelhadas.

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

Em seguida, crie o dataset de entrada e chame tf.distribute.Strategy.experimental_distribute_dataset para distribuir o dataset com base na estratégia.

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

Depois, defina um passo do treinamento. Use tf.GradientTape para computar gradientes e otimizador para aplicar esses gradientes para atualizar as variáveis ​​do seu modelo. Para distribuir este passo de treinamento, coloque-o numa função train_step e passe-a para tf.distribute.Strategy.run junto com as entradas do dataset que você obteve do dist_dataset criado 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)

Algumas outras coisas a serem observadas no código acima:

1. Você usou tf.nn.compute_average_loss para reduzir as perdas de previsão por exemplo para um escalar. tf.nn.compute_average_loss soma a perda por exemplo e divide a soma pelo tamanho global do lote. Isto é importante porque mais tarde, depois que os gradientes forem calculados em cada réplica, eles serão agregados nas réplicas, através da soma.

Por padrão, o tamanho do lote global é tf.get_strategy().num_replicas_in_sync * tf.shape(per_example_loss)[0]. Ele também pode ser especificado explicitamente como um argumento de palavra-chave global_batch_size=. Sem lotes pequenos, o padrão é equivalente a tf.nn.compute_average_loss(..., global_batch_size=global_batch_size) com o global_batch_size definido acima. (Para saber mais sobre lotes pequenos e como evitá-los ou lidar com eles, veja o Tutorial de treinamento personalizado.)

1. Você usou tf.nn.scale_regularization_loss para escalonar as perdas de regularização registradas com o objeto Model, se houver, também em 1/num_replicas_in_sync. Para as perdas de regularização que dependem da entrada, cabe ao código de modelagem, e não ao loop de treinamento personalizado, realizar a média do tamanho do lote por réplica(!). Assim, o código de modelagem pode permanecer agnóstico em relação à replicação, enquanto o loop de treinamento permanecerá agnóstico em relação a forma como as perdas de regularização são computadas.

2. Quando você chama apply_gradients dentro de um escopo de estratégia de distribuição, seu comportamento é modificado. Especificamente, antes de aplicar gradientes em cada instância paralela durante o treinamento síncrono, será realizada a soma de todas as réplicas dos gradientes.

3. Você também usou a API tf.distribute.Strategy.reduce para agregar os resultados retornados por tf.distribute.Strategy.run para a geração de relatórios. tf.distribute.Strategy.run retorna resultados de cada réplica local na estratégia e há várias maneiras de consumir esse resultado. Você pode aplicar reduce neles para obter um valor agregado. Você também pode fazer tf.distribute.Strategy.experimental_local_results para obter a lista de valores contidos no resultado, um por réplica local.

Finalmente, uma vez definido o passo de treinamento, você pode iterar por dist_dataset e executar o treinamento num loop:

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

No exemplo acima, você iterou pelo dist_dataset para fornecer entradas para seu treinamento. Você também recebe o tf.distribute.Strategy.make_experimental_numpy_dataset para garantir suporte a entradas NumPy. Você pode usar esta API para criar um dataset antes de chamar tf.distribute.Strategy.experimental_distribute_dataset.

Outra forma de iterar seus dados é usar iteradores explicitamente. Talvez você queira fazer isso quando precisar executar um determinado número de passos, em vez de iterar por todo o dataset. A iteração acima agora seria modificada para primeiro criar um iterador e, em seguida, chamar next explicitamente para obter os dados de entrada.

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

Isto cobre o caso de uso mais simples da API tf.distribute.Strategy para distribuir loops de treinamento personalizados.

O que suportado agora?

API de treinamento	MirroredStrategy	TPUStrategy	MultiWorkerMirroredStrategy	ParameterServerStrategy	CentralStorageStrategy
Loop de treinamento personalizado	Suportado	Suportado	Suportado	Suporte experimental	Suporte experimental

Exemplos e tutoriais

Aqui estão alguns exemplos de uso de estratégias de distribuição com ciclos de treinamento personalizados:

1. Tutorial: Treinamento com um loop de treinamento personalizado e MirroredStrategy.

2. Tutorial: Treinamento com um loop de treinamento personalizado e MultiWorkerMirroredStrategy.

3. Guia: Contém um exemplo de loop de treinamento personalizado com TPUStrategy.

4. Tutorial: Treinamento de servidor de parâmetros com um loop de treinamento personalizado e ParameterServerStrategy.

5. Repositório TensorFlow Model Garden contendo coleções de modelos de última geração implementados usando diversas estratégias.

Outros tópicos

Esta seção cobre alguns tópicos que são relevantes para diversos casos de uso.

Configurando a variável de ambiente TF_CONFIG

Para treinamento multiworker, conforme mencionado anteriormente, você precisa configurar a variável de ambiente 'TF_CONFIG' para cada binário em execução no seu cluster. A variável de ambiente 'TF_CONFIG' é uma string JSON que especifica quais tarefas constituem um cluster, seus endereços e a função de cada tarefa no cluster. O repositório tensorflow/ecosystem fornece um modelo Kubernetes, que configura 'TF_CONFIG' para suas tarefas de treinamento.

Existem dois componentes de 'TF_CONFIG': um cluster e uma tarefa.

• Um cluster fornece informações sobre o cluster de treinamento, que é um dict que consiste de diferentes tipos de jobs, tais como workers. No treinamento com múltiplos workers, geralmente há um worker que assume um pouco mais de responsabilidade, como salvar checkpoints e escrever um arquivo de resumo para o TensorBoard, além do que um worker normal já faz. Esse worker é chamado de worker "chief", e normalmente o worker com índice 0 é nomeado como sendo o worker chief (na verdade, é dessa forma que tf.distribute.Strategy é implementada).

• Uma tarefa, por outro lado, fornece informações sobre a tarefa atual. O primeiro cluster de componentes é o mesmo para todos os workers, e a tarefa do segundo componente é diferente em cada worker e especifica o tipo e o índice desse worker.

Um exemplo de 'TF_CONFIG' é:

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 existem três workers e duas tarefas "ps" no "cluster" junto com seus hosts e portas. A parte "task" especifica a função da tarefa atual no "cluster" —worker 1 (o segundo worker). As funções válidas em um cluster são "chief", "worker", "ps" e "evaluator". Não deve haver nenhuma tarefa "ps", exceto quando usar tf.distribute.experimental.ParameterServerStrategy.

Quais são os próximos passos?

tf.distribute.Strategy está ativamente em desenvolvimento. Experimente e forneça seus comentários usando issues do GitHub.