Treinamento em pods da TPU

Visão geral

Um dos benefícios da Cloud TPU é a possibilidade de escalonar o treinamento entre tipos de TPU com tamanhos diferentes. Todos os tipos de TPU usam a mesma arquitetura com paralelismo de dados. A única mudança é que o paralelismo aumenta de 8 para 2.048 núcleos.

Para aproveitar ao máximo os números maiores de TPUs, é necessário ajustar vários parâmetros de tarefas de treinamento. Neste documento, explicamos alguns problemas comuns e as alterações que você precisa fazer nos seus modelos.

Como escalonar o tamanho do lote e as etapas de treinamento

Para alcançar o dimensionamento linear em tipos de TPU maiores, mantenha o tamanho do lote por núcleo igual.

Por exemplo, se você usar um tamanho de lote de 1.024 em uma v2-8, use um tamanho de lote de 4.096 (4 * 1024) em uma v2-32. Isso utiliza totalmente o hardware da TPU. É possível usar tamanhos de lote menores, mas o treinamento não será dimensionado linearmente se você fizer isso.

Muitos modelos incluem uma sinalização train_steps em que uma etapa corresponde a um único lote de dados. Quando aumentar o tamanho do lote, reduza o número de etapas de treinamento para que o número total de exemplos de treinamento permaneça o mesmo.

Por exemplo, se a configuração original era de 1024 para 1.000 etapas, o tamanho do lote poderia ser aumentado para 4096 e o número de etapas poderia ser reduzido para 250 (1.000 / 4). Caso o modelo use uma sinalização epochs, não é necessário escalonar o número de etapas.

Tamanhos de lote maiores podem alterar o comportamento de convergência do modelo. Portanto, talvez você também precise ajustar alguns hiperparâmetros.

Como usar otimizadores de lotes grandes

Se você realizar o escalonamento vertical para usar tamanhos de lote maiores, talvez precise usar um otimizador compatível com grandes tamanhos de lote. Por exemplo, o modelo de referência Resnet usa o LARSOptimizer (ambos em inglês) para tipos de TPU maiores.

Como usar buckets regionais do Google Cloud Storage na mesma região do pod de TPU

Em geral, a prática recomendada para treinamento de TPU é sempre usar recursos na mesma região. A região de recursos é importante principalmente ao usar pods de TPU porque a taxa de transferência de dados do Google Cloud Storage tende a ser maior. Verifique se você está usando um bucket regional do Google Cloud Storage na mesma região da TPU para treinar conjuntos de dados e checkpoints.

Como fornecer o nome e a região da TPU para TPUClusterResolver

Para o treinamento em um dispositivo único, especifique o nome da TPU ou um endereço IP, por exemplo: grpc://1.2.3.4:8470.

Para pods de TPU, use o nome da TPU para que o TensorFlow descubra os endereços IP de todos os hosts disponíveis para distribuição de treinamento.

A especificação do endereço IP de um pod de TPU gera o seguinte erro:

ValueError: TPUConfig.num_shards is not set correctly ....

Se for preciso executar o modelo e a TPU em zonas ou regiões diferentes, forneça o TPUClusterResolver com a zona de TPU.

É possível transmitir as sinalizações tpu e tpu_zone para modelos de referência ou definir diretamente os valores para TPUClusterResolver. Por exemplo:

cluster = tf.contrib.cluster_resolver.TPUClusterResolver(
  tpu='tpuname',
  zone='europe-west4-a',
  project='myproject')

Um problema relacionado que pode ocorrer ao usar a TPUEstimator não está passando no resolvedor de cluster para tf.contrib.tpu.RunConfig. Isso pode acontecer se RunConfig for inicializado com o parâmetro master. Em vez disso, inicialize-o com o parâmetro cluster.

config = tf.contrib.tpu.RunConfig(cluster=cluster)

Se você estiver usando um pod de TPU e tf.Session sem usar a TPUEstimator, ajuste a configuração da sessão da seguinte maneira:

with tf.Session(cluster.get_master(), config=config.session_config) as sess:

Avaliação em pods de TPU

A avaliação não é compatível nem econômica em pods de TPU. O Google recomenda executar a avaliação em uma TPU de dispositivo único usando a sinalização mode, compatível com a maioria dos modelos.

Para o treinamento em um pod de TPU ou em uma TPU de dispositivo único, defina a sinalização como train. Para avaliação, defina eval apenas para TPUs de dispositivo único. Também é possível usar o modo train_and_eval em TPUs de dispositivo único, mas não em pods de TPU.