Ter uma questão? Conecte-se com a comunidade no Fórum TensorFlow Visite o Fórum

Treinamento distribuído com TensorFlow

Ver no TensorFlow.org Executar no Google Colab Ver fonte no GitHub Baixar caderno

Visão geral

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

tf.distribute.Strategy foi projetado com estes objetivos principais em mente:

  • Fácil de usar e oferecer suporte a vários segmentos de usuários, incluindo pesquisadores, engenheiros de ML, etc.
  • Fornece bom desempenho fora da caixa.
  • Troca fácil entre estratégias.

tf.distribute.Strategy pode ser usado com uma API de alto nível como Keras e também pode ser usado para distribuir loops de treinamento personalizados (e, em geral, qualquer computação usando TensorFlow).

No TensorFlow 2.x, você pode executar seus programas avidamente ou em um gráfico 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 TPUStrategy . Embora o treinamento seja o foco deste guia, esta API também pode ser usada para distribuir avaliação e predição em diferentes plataformas.

Você pode usar tf.distribute.Strategy com muito poucas alterações em seu código, porque os componentes subjacentes do TensorFlow foram alterados para se tornarem cientes da estratégia. Isso inclui variáveis, camadas, modelos, otimizadores, métricas, resumos e pontos de verificação.

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

# Import TensorFlow
import tensorflow as tf

Tipos de estratégias

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

  • Treinamento síncrono versus assíncrono: essas são duas maneiras comuns de distribuir treinamento com paralelismo de dados. No treinamento de sincronização, todos os funcionários treinam em diferentes fatias de dados de entrada em sincronia e agregam gradientes em cada etapa. No treinamento assíncrono, todos os trabalhadores treinam independentemente sobre os dados de entrada e atualizam as variáveis ​​de forma assíncrona. Normalmente, o treinamento de sincronização é suportado por meio de redução total e assíncrono por meio de arquitetura de servidor de parâmetro.
  • Plataforma de hardware: você pode querer escalar seu treinamento em várias GPUs em uma máquina, ou várias máquinas em uma rede (com 0 ou mais GPUs cada) ou em Cloud TPUs.

Para oferecer suporte a esses casos de uso, existem 6 estratégias disponíveis. A próxima seção explica quais deles são compatíveis em quais cenários no TensorFlow. Aqui está uma visão geral rápida:

API de treinamento MirroredStrategy TPUStrategy MultiWorkerMirroredStrategy CentralStorageStrategy ParameterServerStrategy
API Keras Suportado Suportado Suportado Suporte Experimental Postagem planejada apoiada 2.4
Loop de treinamento personalizado Suportado Suportado Suportado Suporte Experimental Suporte Experimental
API Estimator Suporte Limitado Não suportado Suporte Limitado Suporte Limitado Suporte Limitado

MirroredStrategy

tf.distribute.MirroredStrategy oferece suporte a treinamento distribuído síncrono em várias GPUs em uma máquina. Ele 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 em sincronia entre si aplicando atualizações idênticas.

Algoritmos eficientes de redução total são usados ​​para comunicar as atualizações de variáveis ​​entre os dispositivos. O All-reduz agrega tensores em todos os dispositivos, adicionando-os e tornando-os disponíveis em cada dispositivo. É um algoritmo fundido muito eficiente e pode reduzir significativamente a sobrecarga de sincronização. Existem muitos algoritmos de redução total e implementações 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 a implementação totalmente reduzida . Você pode escolher entre algumas outras opções ou escrever a sua própria.

Esta é a maneira mais simples de criar MirroredStrategy :

mirrored_strategy = tf.distribute.MirroredStrategy()
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)

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

Se desejar usar apenas algumas das GPUs em sua máquina, você pode fazer assim:

mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])
WARNING:tensorflow:Some requested devices in `tf.distribute.Strategy` are not visible to TensorFlow: /job:localhost/replica:0/task:0/device:GPU:1,/job:localhost/replica:0/task:0/device:GPU:0
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1')

Se desejar 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())
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)

TPUStrategy

tf.distribute.TPUStrategy permite que você execute seu treinamento TensorFlow em unidades de processamento de tensor (TPUs). TPUs são ASICs especializados do Google, projetados para acelerar drasticamente as cargas de trabalho de aprendizado de máquina. Eles estão disponíveis no Google Colab , TensorFlow Research Cloud e Cloud TPU .

Em termos de arquitetura de treinamento distribuída, TPUStrategy é a mesma MirroredStrategy - 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 vários núcleos de TPU, que são usados ​​na TPUStrategy .

Aqui está como você instanciaria TPUStrategy :

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.TPUStrategy(cluster_resolver)

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

Se você quiser usar isso para Cloud TPUs:

  • Você deve especificar o nome do seu recurso TPU no argumento tpu .
  • Você deve inicializar o sistema tpu explicitamente no início do programa. Isso é necessário para 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 . Ele implementa treinamento distribuído síncrono entre vários funcionários, 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 trabalhadores.

Esta é a maneira mais simples de criar MultiWorkerMirroredStrategy :

strategy = tf.distribute.MultiWorkerMirroredStrategy()
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:GPU:0',), communication = CommunicationImplementation.AUTO

MultiWorkerMirroredStrategy tem 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 transfere a escolha para o Tensorflow. Você pode especificá-los da seguinte maneira:

communication_options = tf.distribute.experimental.CommunicationOptions(
    implementation=tf.distribute.experimental.CommunicationImplementation.NCCL)
strategy = tf.distribute.MultiWorkerMirroredStrategy(
    communication_options=communication_options)
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:GPU:0',), communication = CommunicationImplementation.NCCL

Uma das principais diferenças para iniciar o treinamento de vários trabalhadores, em comparação com o treinamento de várias GPU, é a configuração de vários trabalhadores. A variável de ambiente TF_CONFIG é a maneira padrão no TensorFlow de especificar a configuração do cluster para cada trabalhador que faz parte do cluster. Saiba mais sobre como configurar TF_CONFIG .

ParameterServerStrategy

O treinamento do servidor de parâmetros é um método comum de paralelismo de dados para dimensionar o treinamento do modelo em várias máquinas. Um cluster de treinamento de servidor de parâmetro consiste em trabalhadores e servidores de parâmetro. As variáveis ​​são criadas em servidores de parâmetros e são lidas e atualizadas pelos trabalhadores em cada etapa. Confira o tutorial de treinamento do servidor de parâmetros para obter detalhes.

No TensorFlow 2, o treinamento do servidor de parâmetros usa uma arquitetura baseada no 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 s que tf.distribute.Server as tarefas do coordenador. O coordenador cria recursos, despacha tarefas de treinamento, escreve pontos de verificação e lida com falhas de tarefas.

Na programação em execução no coordenador, você usará um objeto ParameterServerStrategy para definir uma etapa de treinamento e usará um ClusterCoordinator para despachar as etapas de treinamento para trabalhadores remotos. Esta é 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)

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

CentralStorageStrategy

tf.distribute.experimental.CentralStorageStrategy 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 de CentralStorageStrategy por:

central_storage_strategy = tf.distribute.experimental.CentralStorageStrategy()
INFO:tensorflow:ParameterServerStrategy (CentralStorageStrategy if you are using a single machine) with compute_devices = ['/job:localhost/replica:0/task:0/device:GPU:0'], variable_device = '/job:localhost/replica:0/task:0/device:GPU:0'

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

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 .

Estratégia Padrão

A 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. Ele implementa a interface tf.distribute.Strategy , mas é uma passagem e não fornece distribuição real. Por exemplo, strategy.run(fn) simplesmente chamará fn . O código escrito usando essa estratégia deve se comportar exatamente como o código escrito sem nenhuma estratégia. Você pode pensar nisso como uma estratégia "sem operação".

A estratégia padrão é um singleton - e não se pode criar mais instâncias dela. Ele pode ser obtido usando tf.distribute.get_strategy fora do escopo de qualquer estratégia explícita (a mesma API que pode ser usada para obter a estratégia atual dentro do escopo de uma estratégia explícita).

default_strategy = tf.distribute.get_strategy()

Essa estratégia serve a dois propósitos principais:

  • Ele permite escrever código de biblioteca com reconhecimento de distribuição incondicionalmente. Por exemplo, em tf.optimizer s pode usar tf.distribute.get_strategy e usar essa estratégia para reduzir gradientes - ele sempre retornará um objeto de estratégia no qual você pode chamar a API de redução.
# In optimizer or other library code
# Get currently active strategy
strategy = tf.distribute.get_strategy()
strategy.reduce("SUM", 1., axis=None)  # reduce some values
1.0
  • Semelhante ao código de biblioteca, ele pode ser usado para escrever programas de usuários finais para trabalhar com e sem estratégia de distribuição, sem exigir lógica condicional. Aqui está um exemplo de snippet de código 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.))
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)
MirroredVariable:{
  0: <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>
}

OneDeviceStrategy

tf.distribute.OneDeviceStrategy é uma estratégia para colocar todas as variáveis ​​e cálculos em um único dispositivo especificado.

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

Essa estratégia é diferente da estratégia padrão de várias maneiras. Na estratégia padrão, a lógica de colocação de variável permanece inalterada quando comparada à execução do TensorFlow sem qualquer estratégia de distribuição. Mas ao usar OneDeviceStrategy , todas as variáveis ​​criadas em seu escopo são colocadas explicitamente no dispositivo especificado. Além disso, todas as funções chamadas por meio de OneDeviceStrategy.run também serão colocadas no dispositivo especificado.

A entrada distribuída por meio dessa estratégia será pré-buscada para o dispositivo especificado. Na estratégia padrão, não há distribuição de entrada.

Semelhante à estratégia padrão, esta estratégia também pode ser usada para testar seu código antes de mudar para outras estratégias que realmente são distribuídas para vários dispositivos / máquinas. Isso exercitará a maquinaria da estratégia de distribuição um pouco mais do que a estratégia padrão, mas não em toda a extensão do uso de MirroredStrategy ou TPUStrategy etc. Se você quiser um código que se comporte como se não tivesse estratégia, use a estratégia padrão.

Até agora, você viu as diferentes estratégias disponíveis e como pode instanciá-las. As próximas seções mostram as diferentes maneiras em que você pode usá-los para distribuir seu treinamento.

Usando tf.distribute.Strategy com tf.keras.Model.fit

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

Aqui está 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 .

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 densa:

mirrored_strategy = tf.distribute.MirroredStrategy()

with mirrored_strategy.scope():
  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])

model.compile(loss='mse', optimizer='sgd')
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).

Este exemplo usa MirroredStrategy , portanto, você pode executá-lo em uma máquina com várias GPUs. strategy.scope() indica a Keras qual estratégia usar para distribuir o treinamento. A criação de modelos / otimizadores / métricas dentro deste escopo permite que você crie variáveis ​​distribuídas em vez de variáveis ​​regulares. Assim que estiver 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)
Epoch 1/2
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
10/10 [==============================] - 3s 2ms/step - loss: 0.3648
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
Epoch 2/2
10/10 [==============================] - 0s 2ms/step - loss: 0.1613
10/10 [==============================] - 1s 2ms/step - loss: 0.1001
0.1001412644982338

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

import numpy as np

inputs, targets = np.ones((100, 1)), np.ones((100, 1))
model.fit(inputs, targets, epochs=2, batch_size=10)
Epoch 1/2
10/10 [==============================] - 1s 2ms/step - loss: 0.0713
Epoch 2/2
10/10 [==============================] - 0s 2ms/step - loss: 0.0315
<tensorflow.python.keras.callbacks.History at 0x7f18ac0bdcd0>

Em ambos os casos - com Dataset ou Numpy - cada lote da entrada fornecida é dividido igualmente entre as várias réplicas. Por exemplo, se você estiver usando a 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 etapa. Cada época será treinada mais rápido conforme você adiciona mais GPUs. Normalmente, você deseja aumentar o tamanho do lote à medida que adiciona mais aceleradores, para fazer uso eficaz do poder de computação extra. 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.

# Compute global batch size using 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}
learning_rate = LEARNING_RATES_BY_BATCH_SIZE[global_batch_size]

O que é compatível agora?

API de treinamento MirroredStrategy TPUStrategy MultiWorkerMirroredStrategy ParameterServerStrategy CentralStorageStrategy
APIs Keras Suportado Suportado Suporte Experimental 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 Keras:

  1. Tutorial para treinar MNIST com MirroredStrategy .
  2. Tutorial para treinar MNIST usando MultiWorkerMirroredStrategy .
  3. Guia de treinamento MNIST usando TPUStrategy .
  4. Tutorial para treinamento de servidor de parâmetro no TensorFlow 2 com ParameterServerStrategy .
  5. TensorFlow Modelo Garden repositório contendo coleções de modelos state-of-the-art implementadas usando várias estratégias.

Usando tf.distribute.Strategy com loops de treinamento personalizados

Como você viu, usar tf.distribute.Strategy com Keras model.fit requer a alteração de apenas algumas linhas de 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ê precisa de mais flexibilidade e controle sobre seus loops de treinamento do que é possível com o Estimator ou Keras, você pode escrever loops de treinamento personalizados. Por exemplo, ao usar um GAN, você pode querer realizar um número diferente de etapas de gerador ou discriminador a cada rodada. Da mesma forma, as estruturas de alto nível não são muito adequadas 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. O uso deles pode exigir uma pequena reestruturação do código inicialmente, mas depois que isso for feito, você poderá alternar entre GPUs, TPUs e várias máquinas simplesmente alterando a instância da estratégia.

Aqui, você verá um breve trecho ilustrando 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 garante que quaisquer variáveis ​​criadas com o modelo e otimizador sejam variáveis ​​espelhadas.

with mirrored_strategy.scope():
  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])
  optimizer = tf.keras.optimizers.SGD()

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

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

Em seguida, defina uma etapa do treinamento. Usetf.GradientTape para calcular gradientes e otimizador para aplicar esses gradientes para atualizar as variáveis ​​do seu modelo. Para distribuir esta etapa de treinamento, coloque-a em uma função train_step e passe-a para tf.distrbute.Strategy.run junto com as entradas do conjunto de dados que você obteve do dist_dataset criado antes:

loss_object = tf.keras.losses.BinaryCrossentropy(
  from_logits=True,
  reduction=tf.keras.losses.Reduction.NONE)

def compute_loss(labels, predictions):
  per_example_loss = loss_object(labels, predictions)
  return tf.nn.compute_average_loss(per_example_loss, global_batch_size=global_batch_size)

def train_step(inputs):
  features, labels = inputs

  with tf.GradientTape() as tape:
    predictions = model(features, training=True)
    loss = compute_loss(labels, predictions)

  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 calcular a perda. tf.nn.compute_average_loss soma a perda por exemplo e divide a soma pelo global_batch_size. Isso é importante porque, mais tarde, depois que os gradientes são calculados em cada réplica, eles são agregados entre as réplicas somando- os.
  2. Você também usou a API tf.distribute.Strategy.reduce para agregar os resultados retornados por tf.distribute.Strategy.run . 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 reduce los 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.
  3. 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, ele executa uma soma de todas as réplicas dos gradientes.

Finalmente, depois de definir a etapa de treinamento, você pode iterar em dist_dataset e executar o treinamento em um loop:

for dist_inputs in dist_dataset:
  print(distributed_train_step(dist_inputs))
tf.Tensor(0.5716536, shape=(), dtype=float32)
tf.Tensor(0.56787133, shape=(), dtype=float32)
tf.Tensor(0.564126, shape=(), dtype=float32)
tf.Tensor(0.5604175, shape=(), dtype=float32)
tf.Tensor(0.5567452, shape=(), dtype=float32)
tf.Tensor(0.5531087, shape=(), dtype=float32)
tf.Tensor(0.54950774, shape=(), dtype=float32)
tf.Tensor(0.5459418, shape=(), dtype=float32)
tf.Tensor(0.5424106, shape=(), dtype=float32)
tf.Tensor(0.5389137, shape=(), dtype=float32)
tf.Tensor(0.53545076, shape=(), dtype=float32)
tf.Tensor(0.5320213, shape=(), dtype=float32)
tf.Tensor(0.528625, shape=(), dtype=float32)
tf.Tensor(0.5252616, shape=(), dtype=float32)
tf.Tensor(0.52193063, shape=(), dtype=float32)
tf.Tensor(0.5186317, shape=(), dtype=float32)
tf.Tensor(0.5153646, shape=(), dtype=float32)
tf.Tensor(0.51212865, shape=(), dtype=float32)
tf.Tensor(0.5089238, shape=(), dtype=float32)
tf.Tensor(0.5057497, shape=(), dtype=float32)

No exemplo acima, você dist_dataset no dist_dataset para fornecer entrada para seu treinamento. Você também tf.distribute.Strategy.make_experimental_numpy_dataset o tf.distribute.Strategy.make_experimental_numpy_dataset para suportar entradas NumPy. Você pode usar esta API para criar um conjunto de dados antes de chamar tf.distribute.Strategy.experimental_distribute_dataset .

Outra maneira de iterar seus dados é usar explicitamente iteradores. Você pode querer fazer isso quando quiser executar um determinado número de etapas, em vez de iterar por todo o conjunto de dados. A iteração acima agora seria modificada para criar primeiro um iterador e, em seguida, chamar explicitamente next nele para obter os dados de entrada.

iterator = iter(dist_dataset)
for _ in range(10):
  print(distributed_train_step(next(iterator)))
tf.Tensor(0.5026059, shape=(), dtype=float32)
tf.Tensor(0.49949214, shape=(), dtype=float32)
tf.Tensor(0.496408, shape=(), dtype=float32)
tf.Tensor(0.4933532, shape=(), dtype=float32)
tf.Tensor(0.49032742, shape=(), dtype=float32)
tf.Tensor(0.48733026, shape=(), dtype=float32)
tf.Tensor(0.48436147, shape=(), dtype=float32)
tf.Tensor(0.4814206, shape=(), dtype=float32)
tf.Tensor(0.4785076, shape=(), dtype=float32)
tf.Tensor(0.475622, shape=(), dtype=float32)

Isso cobre o caso mais simples de usar a API tf.distribute.Strategy para distribuir loops de treinamento personalizados.

O que é compatível agora?

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

Exemplos e tutoriais

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

  1. Tutorial para treinar MNIST usando MirroredStrategy .
  2. Guia de treinamento MNIST usando TPUStrategy .
  3. TensorFlow Modelo Garden repositório contendo coleções de modelos state-of-the-art implementadas usando várias estratégias.

Outros tópicos

Esta seção cobre alguns tópicos relevantes para vários casos de uso.

Configurando a variável de ambiente TF_CONFIG

Para treinamento de vários trabalhadores, conforme mencionado antes, você precisa configurar a variável de ambiente TF_CONFIG para cada binário em execução em 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 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 dicionário que consiste em diferentes tipos de trabalhos, como trabalhador. No treinamento de vários trabalhadores, geralmente há um trabalhador que assume um pouco mais de responsabilidade, como salvar o ponto de verificação e escrever um arquivo de resumo para o TensorBoard, além do que um trabalhador regular faz. Esse trabalhador é referido como o trabalhador "chefe" e é comum que o trabalhador com índice 0 seja indicado como o trabalhador chefe (na verdade, é assim que tf.distribute.Strategy é implementado).
  • Por outro lado, uma tarefa fornece informações sobre a tarefa atual. O primeiro cluster de componente é o mesmo para todos os trabalhadores e a segunda tarefa do componente é diferente em cada trabalhador e especifica o tipo e o índice desse trabalhador.

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

Qual é o próximo?

tf.distribute.Strategy está ativamente em desenvolvimento. Experimente e forneça seus comentários sobre os problemas do GitHub .