Visão geral
A versão do software do framework em execução na TPU precisa corresponder à versão em execução na VM local. Agora é possível ativar essa versão de software em um Cloud TPU em execução, sem excluir e recriar a TPU. Isso também permite configurar a Cloud TPU com versões noturnas específicas de frameworks de software. Ainda é recomendável selecionar uma versão compatível com esses frameworks.
Uso
A maneira recomendada de alternar as versões é usar a biblioteca python cloud-tpu-client.
Exemplo de uso do TensorFlow.
import tensorflow as tf
from cloud_tpu_client import Client
c = Client()
c.configure_tpu_version(tf.__version__, restart_type='ifNeeded')
Isso configura a Cloud TPU para corresponder à versão do TensorFlow em execução na sua VM local. Isso inclui versões oficiais e versões noturnas datadas.
O parâmetro restart_type da API configure_tpu_version define o comportamento de reinicialização da TPU ao alternar versões. As opções são
"sempre" (o padrão) e "ifNeeded".
"sempre" pode ser usado para corrigir uma TPU com, por exemplo, status UNHEALTHY_TENSORFLOW ou que esteja retornando erros de falta de memória (OOM, na sigla em inglês) devido a recursos vazados de uma execução anterior. Quando essa opção é definida, a TPU é reiniciada mesmo quando uma nova versão do framework não está instalada.
"ifNeeded" pode ser útil porque não reinicia o ambiente de execução se já estiver na versão correta, por isso não adicionará qualquer tempo de inicialização significativo a um script de treinamento. Quando essa opção é configurada, a TPU é reiniciada apenas se não tiver a versão correta do framework instalada.
A biblioteca se comunica diretamente com a Cloud TPU. Portanto, esse código precisa ser executado em uma VM na mesma rede. É recomendável executar isso no código para o restante do modelo.
Outras opções de software
O TensorFlow inclui uma string tf.__version__ que é a maneira mais simples
de configurar a versão correta. Outras opções de software incluem:
- PyTorch -
pytorch-0.1,pytorch-0.5-dev20200320,pytorch-nightly - Jax -
tpu_driver,tpu_driver0.1-dev20200320,tpu_driver_nightly
Por exemplo, para configurar uma TPU para ser executada com a versão noturna mais recente do PyTorch.
from cloud_tpu_client import Client
c = Client()
c.configure_tpu_version('pytorch-nightly', restart_type='ifNeeded')
c.wait_for_healthy()