Como realizar o treinamento em um pod de 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.

No entanto, é preciso ajustar vários parâmetros ao realizar o treinamento em tipos de TPU maiores. Neste documento, você verá alguns problemas comuns e as modificações necessárias nos 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 permite o uso integral do hardware da TPU. É possível usar tamanhos de lote menores. No entanto, seu treinamento não será escalonado linearmente se você fizer isso.

Muitos modelos incluem uma sinalização train_steps em que uma etapa corresponde a um único lote de dados. Se você aumentar o tamanho do lote, reduza verticalmente o número de etapas de treinamento. Assim, o número total de exemplos de treinamento permanece o mesmo. Por exemplo, se a configuração original tinha um tamanho de lote de 1.024 para 1.000 etapas, depois de aumentar o tamanho do lote para 4.096, reduza as etapas para 250 (1.000 / 4). Se o modelo usar uma sinalização epochs, não será preciso 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 intervalos 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 intervalo 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.