Executar um cálculo em uma VM da Cloud TPU usando o JAX

Criar uma VM do Cloud TPU usando gcloud

1. Defina algumas variáveis de ambiente para facilitar o uso dos comandos.

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=us-east5-a
   export ACCELERATOR_TYPE=v5litepod-8
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite

   Descrições das variáveis de ambiente

   Variável	Descrição
   PROJECT_ID	O ID do seu Google Cloud projeto. Use um projeto existente ou crie um novo.
   TPU_NAME	O nome da TPU.
   ZONE	A zona em que a VM TPU será criada. Para mais informações sobre as zonas com suporte, consulte Regiões e zonas de TPU.
   ACCELERATOR_TYPE	O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores compatíveis com cada versão de TPU, consulte Versões de TPU.
   RUNTIME_VERSION	A versão do software do Cloud TPU.

2. Crie a VM de TPU executando o comando a seguir em um Cloud Shell ou no terminal do computador em que a Google Cloud CLI está instalada.

   $ gcloud compute tpus tpu-vm create $TPU_NAME \
       --project=$PROJECT_ID \
       --zone=$ZONE \
       --accelerator-type=$ACCELERATOR_TYPE \
       --version=$RUNTIME_VERSION

Conecte-se à VM do Cloud TPU

Conecte-se à VM da TPU usando SSH com o seguinte comando:

$ gcloud compute tpus tpu-vm ssh $TPU_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE

Se você não conseguir se conectar a uma VM da TPU usando SSH, talvez seja porque a VM da TPU não tem um endereço IP externo. Para acessar uma VM da TPU sem um endereço IP externo, siga as instruções em Conectar-se a uma VM da TPU sem um endereço IP público.

Instalar o JAX na VM do Cloud TPU

(vm)$ pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

Verificação do sistema

Verifique se o JAX pode acessar o TPU e executar operações básicas:

1. Inicie o interpretador do Python 3:

   (vm)$ python3

   >>> import jax

2. Veja o número de núcleos de TPU disponíveis:

   >>> jax.device_count()

O número de núcleos de TPU é exibido. O número de cores exibidas depende da versão da TPU que você está usando. Para mais informações, consulte Versões da TPU.

Fazer um cálculo

>>> jax.numpy.add(1, 1)

O resultado da adição "numpy" é exibido:

Resultado do comando:

Array(2, dtype=int32, weak_type=True)

Sair do interpretador Python

>>> exit()

Como executar código do JAX em uma VM de TPU

Agora é possível executar qualquer código do JAX que você quiser. Os exemplos de Flax são um ótimo lugar para começar a executar modelos padrão de ML no JAX. Por exemplo, para treinar uma rede convolucional básica MNIST:

1. Instale as dependências dos exemplos do Flax:

   (vm)$ pip install --upgrade clu
   (vm)$ pip install tensorflow
   (vm)$ pip install tensorflow_datasets

2. Instale o FLAX:

   (vm)$ git clone https://github.com/google/flax.git
   (vm)$ pip install --user flax

3. Execute o script de treinamento do Flax MNIST:

   (vm)$ cd flax/examples/mnist
   (vm)$ python3 main.py --workdir=/tmp/mnist \
       --config=configs/default.py \
       --config.learning_rate=0.05 \
       --config.num_epochs=5

O script faz o download do conjunto de dados e inicia o treinamento. A saída do script será semelhante a esta:

I0214 18:00:50.660087 140369022753856 train.py:146] epoch:  1, train_loss: 0.2421, train_accuracy: 92.97, test_loss: 0.0615, test_accuracy: 97.88
I0214 18:00:52.015867 140369022753856 train.py:146] epoch:  2, train_loss: 0.0594, train_accuracy: 98.16, test_loss: 0.0412, test_accuracy: 98.72
I0214 18:00:53.377511 140369022753856 train.py:146] epoch:  3, train_loss: 0.0418, train_accuracy: 98.72, test_loss: 0.0296, test_accuracy: 99.04
I0214 18:00:54.727168 140369022753856 train.py:146] epoch:  4, train_loss: 0.0305, train_accuracy: 99.06, test_loss: 0.0257, test_accuracy: 99.15
I0214 18:00:56.082807 140369022753856 train.py:146] epoch:  5, train_loss: 0.0252, train_accuracy: 99.20, test_loss: 0.0263, test_accuracy: 99.18