Como trocar versões de software na Cloud TPU

Visão geral

Para usar o Cloud TPU, é preciso selecionar uma versão do software que corresponda ao software em execução na VM local. Essa versão do software agora pode ser ativada no Cloud TPU sem excluir e recriar a TPU. Isso também permite configurar o Cloud TPU com versões noturnas específicas de estruturas de software. Ainda é recomendável selecionar uma versão compatível dessas bibliotecas.

Uso

A maneira recomendada de alternar versões é usar a biblioteca em 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 configurará o Cloud TPU para corresponder à versão exata do TensorFlow em execução na VM local, incluindo versões oficiais e versões noturnas obsoletas.

A biblioteca se comunica diretamente com o Cloud TPU para que esse código precise ser executado em uma VM na mesma rede. É recomendável executar isso no código para o restante do modelo.

Opções adicionais 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')