Treinamento do Cloud TPU v5e

O Cloud TPU v5e é o acelerador de IA de última geração do Google Cloud. Com um tamanho menor de 256 chips por pod, o v5e é otimizado para ser o produto de maior valor para treinamento, ajuste, ajuste e disponibilização de transformadores, texto para imagem e rede neural convolucional (CNN, na sigla em inglês). Para mais informações sobre como usar o Cloud TPU v5e para exibição, consulte Inferência usando a v5e.

Para mais informações sobre o hardware e as configurações de TPU do Cloud TPU v5e, consulte TPU v5e.

Começar

As seções a seguir descrevem como começar a usar a TPU v5e.

Solicitação de cotas

É preciso ter cota para usar a TPU v5e no treinamento. Há diferentes tipos de cota para TPUs sob demanda, TPUs reservadas e VMs spot da TPU. Há cotas separadas necessárias se você estiver usando a TPU v5e para inferência. Para mais informações sobre cotas, consulte Cotas. Para solicitar a cota de TPU v5e, entre em contato com a equipe de vendas do Cloud.

Criar um projeto e uma conta do Google Cloud

Você precisa de uma conta e um projeto do Google Cloud para usar o Cloud TPU. Para mais informações, consulte Configurar um ambiente do Cloud TPU.

Criar uma Cloud TPU

A prática recomendada é provisionar o Cloud TPU v5es como recursos na fila usando o comando queued-resource create. Para mais informações, consulte Gerenciar recursos na fila.

Também é possível usar a API Create Node (gcloud compute tpus tpu-vm create) para provisionar o Cloud TPU v5es. Para mais informações, consulte Gerenciar recursos de TPU.

Para mais informações sobre as configurações disponíveis da v5e para treinamento, consulte Tipos da Cloud TPU v5e para treinamento.

Configuração do framework

Nesta seção, descrevemos o processo geral de configuração do treinamento de modelo personalizado usando JAX ou PyTorch com TPU v5e. O suporte ao TensorFlow está disponível nas versões de ambiente de execução da TPU tpu-vm-tf-2.16.1-pjrt e tpu-vm-tf-2.16.1-pod-pjrt.

Para instruções de configuração de inferência, consulte a Introdução à inferência da v5e.

Configuração para JAX

Se você tiver formas de fração com mais de 8 ícones, terá várias VMs em uma fração. Nesse caso, é necessário usar a sinalização --worker=all para executar a instalação em todas as VMs da TPU em uma única etapa, sem usar o SSH para fazer login em cada uma delas separadamente:

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

É possível executar o comando a seguir para verificar o número de dispositivos. As saídas mostradas aqui foram produzidas com uma fração do v5litepod-16. Esse código testa se tudo está instalado corretamente, verificando se o JAX vê os TensorCores do Cloud TPU e pode executar operações básicas:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

A saída será semelhante a esta:

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

jax.device_count() mostra o número total de ícones em uma determinada fração. jax.local_device_count() indica a contagem de ícones acessíveis por uma única VM nessa fração.

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

A saída será semelhante a esta:

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

Veja os Tutoriais do JAX neste documento para começar com o treinamento da v5e usando o JAX.

Configuração do PyTorch

A v5e só oferece suporte ao ambiente de execução PJRT, e o PyTorch 2.1 ou mais recente usará a PJRT como o ambiente de execução padrão para todas as versões de TPU.

Nesta seção, descrevemos como começar a usar o PJRT na v5e com o PyTorch/XLA com comandos para todos os workers.

Instalar dependências

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'

Se você receber um erro ao instalar as rodas para torch, torch_xla ou torchvision como pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222, faça downgrade da versão com este comando:

pip3 install setuptools==62.1.0

Executar um script com PJRT

Observação: para modelos com alocações consideráveis e frequentes, o uso de tcmalloc pode melhorar significativamente o tempo de treinamento em comparação com a implementação padrão de malloc. Portanto, o malloc padrão usado em uma VM de TPU é tcmalloc. No entanto, dependendo da carga de trabalho (por exemplo, com DLRM que tem alocações muito grandes para as tabelas incorporadas), tcmalloc pode causar uma lentidão. Nesse caso, tente cancelar a definição da variável a seguir usando o malloc padrão:

unset LD_PRELOAD

Veja a seguir um exemplo que usa um script Python para fazer um cálculo em uma VM v5e:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker all \
   --command='
      export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.local/lib/
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      unset LD_PRELOAD
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      python3 -c "import torch; import torch_xla; import torch_xla.core.xla_model as xm; print(xm.xla_device()); dev = xm.xla_device(); t1 = torch.randn(3,3,device=dev); t2 = torch.randn(3,3,device=dev); print(t1 + t2)"'

Isso gera um resultado semelhante ao seguinte:

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')

Veja os Tutoriais do PyTorch neste documento para começar com o treinamento v5e usando o PyTorch.

Exclua a TPU e o recurso na fila no final da sessão. Para excluir um recurso na fila, exclua a fração e, em seguida, o recurso na fila em duas etapas:

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Essas duas etapas também podem ser usadas para remover solicitações de recursos na fila que estão no estado FAILED.

Exemplos de JAX/FLAX

As seções a seguir descrevem exemplos de como treinar modelos JAX e FLAX na TPU v5e.

Treinar o ImageNet na v5e

Neste tutorial, descrevemos como treinar o ImageNet na v5e usando dados de entrada falsos. Se você quiser usar dados reais, consulte o arquivo README no GitHub (em inglês).

Configurar

1. Crie variáveis de ambiente:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d
   

2. Crie um recurso de TPU:

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   Observação: a sinalização QUOTA_TYPE pode ser reserved ou best-effort. Se nenhum deles for especificado, o padrão será on-demand. Acesse quotas para mais informações sobre os diferentes tipos de cotas compatíveis com o Cloud TPU.

   Será possível estabelecer uma conexão SSH com sua VM da TPU quando o recurso na fila estiver no estado ACTIVE:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   Quando o QueuedResource estiver no estado ACTIVE, a saída será semelhante a esta:

    state: ACTIVE
   

3. Instale a versão mais recente do JAX e do jaxlib:

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

   Observação: se aparecer um erro sugerindo que você use jax>=0.4.8, ignore a mensagem.

4. Clone o modelo do ImageNet e instale os requisitos correspondentes:

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='git clone https://github.com/google/flax.git && cd flax/examples/imagenet && pip install -r requirements.txt && pip install flax==0.7.4'
   

5. Para gerar dados fictícios, o modelo precisa de informações sobre as dimensões do conjunto de dados. Isso pode ser coletado dos metadados do conjunto de dados do ImageNet:

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} --zone=${ZONE} --worker=all --command='mkdir -p $HOME/flax/.tfds/metadata/imagenet2012/5.1.0 && curl https://raw.githubusercontent.com/tensorflow/datasets/v4.4.0/tensorflow_datasets/testing/metadata/imagenet2012/5.1.0/dataset_info.json --output $HOME/flax/.tfds/metadata/imagenet2012/5.1.0/dataset_info.json'
   

Treinar o modelo

Depois que todas as etapas anteriores forem concluídas, você poderá treinar o modelo.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd flax/examples/imagenet && JAX_PLATFORMS=tpu python3 imagenet_fake_data_benchmark.py'

Excluir a TPU e o recurso na fila

Exclua a TPU e o recurso na fila no final da sessão.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Modelos Hugging Face FLAX

Os modelos Hugging Face implementados no FLAX funcionam imediatamente no Cloud TPU v5e. Nesta seção, fornecemos instruções para executar modelos conhecidos.

Treinar ViT no Imagenette

Neste tutorial, mostramos como treinar o modelo Vision Transformer (ViT) da HuggingFace usando o conjunto de dados Imagenette da Fast AI no Cloud TPU v5e.

O modelo ViT foi o primeiro que treinou com sucesso um codificador de transformador no ImageNet com resultados excelentes em comparação com redes convolucionais. Para mais informações, consulte os recursos a seguir:

• Visão geral do ViT
• Recurso de código ViT

Configurar

1. Crie variáveis de ambiente:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d
   

2. Crie um recurso de TPU:

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   Observação: a sinalização QUOTA_TYPE pode ser reserved ou best-effort. Se nenhum deles for especificado, o padrão será on-demand. Acesse quotas para mais informações sobre os diferentes tipos de cotas compatíveis com o Cloud TPU.

   Será possível estabelecer uma conexão SSH com sua VM da TPU quando o recurso na fila estiver no estado ACTIVE:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   Quando o recurso na fila estiver no estado ACTIVE, a saída será semelhante a esta:

    state: ACTIVE
   

3. Instale o JAX e a biblioteca dele:

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. Faça o download do repositório Hugging Face e instale os requisitos:

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow==2.16.1 && pip install -r examples/flax/vision/requirements.txt'
   

5. Faça o download do conjunto de dados do Imagenette:

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='cd transformers && wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'
   

Treinar o modelo

Treine o modelo com um buffer pré-mapeado de 4 GB.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers && JAX_PLATFORMS=tpu python3 examples/flax/vision/run_image_classification.py --train_dir "imagenette2/train" --validation_dir "imagenette2/val" --output_dir "./vit-imagenette" --learning_rate 1e-3 --preprocessing_num_workers 32 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --model_name_or_path google/vit-base-patch16-224-in21k --num_train_epochs 3'

Excluir a TPU e o recurso na fila

Exclua a TPU e o recurso na fila no fim da sessão.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Resultados do comparativo de mercado de viT

O script de treinamento foi executado nas versões v5litepod-4, v5litepod-16 e v5litepod-64. A tabela a seguir mostra as capacidades com diferentes tipos de acelerador.

Tipo de acelerador	v5litepod-4	v5litepod-16	v5litepod-64
Época	3	3	3
Tamanho global do lote	32	128	512
Capacidade de processamento (exemplos/segundo)	263,40	429,34	470,71

Treine o Diffusion em Pokémon

Neste tutorial, mostramos como treinar o modelo de difusão estável da HuggingFace usando o conjunto de dados Pokémon no Cloud TPU v5e.

O modelo de difusão estável é um modelo latente de texto para imagem que gera imagens fotorrealistas de qualquer entrada de texto. Para saber mais, acesse os recursos a seguir (links em inglês):

• Visão geral da Difusão estável
• Fonte do código de difusão estável

Configurar

1. Crie variáveis de ambiente:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d
   

2. Crie um recurso de TPU:

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   Observação: a sinalização QUOTA_TYPE pode ser reserved ou best-effort. Se nenhum deles for especificado, o padrão será on-demand. Acesse quotas para mais informações sobre os diferentes tipos de cotas compatíveis com o Cloud TPU.

   Será possível estabelecer uma conexão SSH com sua VM da TPU quando o recurso na fila estiver no estado ACTIVE:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   Quando o recurso na fila estiver no estado ACTIVE, a saída será semelhante a esta:

    state: ACTIVE
   

3. Instale o JAX e a biblioteca dele.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. Faça o download do repositório HuggingFace e instale os requisitos.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='git clone https://github.com/RissyRan/diffusers.git && cd diffusers && pip install . && pip install tensorflow==2.16.1 clu && pip install -U -r examples/text_to_image/requirements_flax.txt'
   

Treinar o modelo

Treine o modelo com um buffer pré-mapeado de 4 GB.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --dataset_name=lambdalabs/pokemon-blip-captions --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model'

Excluir a TPU e o recurso na fila

Exclua a TPU e o recurso na fila no final da sessão.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Resultados comparativos de difusão

O script de treinamento foi executado nas versões v5litepod-4, v5litepod-16 e v5litepod-64. A tabela a seguir mostra as capacidades de processamento.

Tipo de acelerador	v5litepod-4	v5litepod-16	v5litepod-64
Degrau de trem	1.500	1.500	1.500
Tamanho global do lote	32	64	128
Capacidade de processamento (exemplos/segundo)	36,53	43,71	49,36

Treinar a GPT2 no conjunto de dados do OSCAR

Neste tutorial, mostramos como treinar o modelo GPT2 da HuggingFace usando o conjunto de dados OSCAR no Cloud TPU v5e.

O GPT2 é um modelo de transformador pré-treinado em textos brutos sem rotulação humana. Ele foi treinado para prever a próxima palavra em frases. Para saber mais, acesse os recursos a seguir:

• Informações gerais sobre a GPT2
• Fonte do código da GPT2

Configurar

1. Crie variáveis de ambiente:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d
   

2. Crie um recurso de TPU:

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   Observação: a sinalização QUOTA_TYPE pode ser reserved ou best-effort. Se nenhum deles for especificado, o padrão será on-demand. Acesse quotas para mais informações sobre os diferentes tipos de cotas compatíveis com o Cloud TPU.

   Será possível estabelecer uma conexão SSH com sua VM da TPU quando o recurso na fila estiver no estado ACTIVE:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   Quando o recurso na fila estiver no estado ACTIVE, a saída será semelhante a esta:

    state: ACTIVE
   

3. Instale o JAX e a biblioteca dele.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. Faça o download do repositório HuggingFace e instale os requisitos.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow && pip install -r examples/flax/language-modeling/requirements.txt'
   

5. Faça o download das configurações para treinar o modelo.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='cd transformers/examples/flax/language-modeling && gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt .'
   

Treinar o modelo

Treine o modelo com um buffer pré-mapeado de 4 GB.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers/examples/flax/language-modeling && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu python3 run_clm_flax.py --output_dir=./gpt --model_type=gpt2 --config_name=./gpt --tokenizer_name=./gpt --dataset_name=oscar --dataset_config_name=unshuffled_deduplicated_no --do_train --do_eval --block_size=512 --per_device_train_batch_size=4 --per_device_eval_batch_size=4 --learning_rate=5e-3 --warmup_steps=1000 --adam_beta1=0.9 --adam_beta2=0.98 --weight_decay=0.01 --overwrite_output_dir --num_train_epochs=3 --logging_steps=500 --eval_steps=2500'

Excluir a TPU e o recurso na fila

Exclua a TPU e o recurso na fila no final da sessão.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Resultados do comparativo de mercado para GPT2

O script de treinamento foi executado nas versões v5litepod-4, v5litepod-16 e v5litepod-64. A tabela a seguir mostra as capacidades de processamento.

	v5litepod-4	v5litepod-16	v5litepod-64
Época	3	3	3
Tamanho global do lote	64	64	64
Capacidade de processamento (exemplos/segundo)	74,60	72,97	72,62

PyTorch/XLA

As seções a seguir descrevem exemplos de como treinar modelos PyTorch/XLA na TPU v5e.

Treinar o ResNet usando o ambiente de execução PJRT

O PyTorch/XLA está migrando de XRT para PjRt do PyTorch 2.0+. Confira as instruções atualizadas para configurar a v5e para cargas de trabalho de treinamento PyTorch/XLA.

Configurar

1. Crie variáveis de ambiente:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=tpu-name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d
   

2. Crie um recurso de TPU:

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --{QUOTA_TYPE}
   

   Observação: a sinalização QUOTA_TYPE pode ser reserved ou best-effort. Se nenhum deles for especificado, o padrão será on-demand. Acesse quotas para mais informações sobre os diferentes tipos de cotas compatíveis com o Cloud TPU.

   Será possível estabelecer uma conexão SSH com sua VM da TPU quando o QueuedResource estiver no estado ACTIVE:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   Quando o recurso na fila estiver no estado ACTIVE, a saída será semelhante a esta:

    state: ACTIVE
   

3. Instalar as dependências específicas do Torch/XLA

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='
         sudo apt-get update -y
         sudo apt-get install libomp5 -y
         pip3 install mkl mkl-include
         pip3 install tf-nightly tb-nightly tbp-nightly
         pip3 install numpy
         sudo apt-get install libopenblas-dev -y
         pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
         pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
   

Treinar o modelo ResNet

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      date
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export XLA_USE_BF16=1
      export LIBTPU_INIT_ARGS=--xla_jf_auto_cross_replica_sharding
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      git clone https://github.com/pytorch/xla.git
      cd xla/
      git reset --hard caf5168785c081cd7eb60b49fe4fffeb894c39d9
      python3 test/test_train_mp_imagenet.py --model=resnet50  --fake_data --num_epochs=1 —num_workers=16  --log_steps=300 --batch_size=64 --profile'

Excluir a TPU e o recurso na fila

Exclua a TPU e o recurso na fila no final da sessão.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Resultado do comparativo de mercado

A tabela a seguir mostra as capacidades de processamento do comparativo.

Tipo de acelerador	Capacidade de processamento (exemplos/segundo)
v5litepod-4	4240 ex/s
v5litepod-16	10.810 ex/s
v5litepod-64	46.154 ex/s

Treinar o GPT2 na v5e

Este tutorial mostra como executar o GPT2 na v5e usando o repositório HuggingFace em PyTorch/XLA usando o conjunto de dados de wikitext.

Configurar

1. Crie variáveis de ambiente:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d
   

2. Crie um recurso de TPU:

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   Observação: a sinalização QUOTA_TYPE pode ser reserved ou best-effort. Se nenhum deles for especificado, o padrão será on-demand. Acesse quotas para mais informações sobre os diferentes tipos de cotas compatíveis com o Cloud TPU.

   Será possível estabelecer uma conexão SSH com sua VM da TPU quando o QueuedResource estiver no estado ACTIVE:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   Quando o recurso na fila está no estado ACTIVE, a saída é semelhante a esta:

   state: ACTIVE
   

3. Instale as dependências do PyTorch/XLA.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='
         sudo apt-get -y update
         sudo apt install -y libopenblas-base
         pip3 install torchvision
         pip3 uninstall -y torch
         pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
         pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
   

4. Faça o download do repositório HuggingFace e instale os requisitos.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='
         git clone https://github.com/pytorch/xla.git
         pip install --upgrade accelerate
         git clone https://github.com/huggingface/transformers.git
         cd transformers
         git checkout ebdb185befaa821304d461ed6aa20a17e4dc3aa2
         pip install .
         git log -1
         pip install datasets evaluate scikit-learn
         '
   

5. Fazer o download das configurações do modelo pré-treinado.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='
         gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/my_config_2.json transformers/examples/pytorch/language-modeling/
         gsutil cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/fsdp_config.json transformers/examples/pytorch/language-modeling/'
   

Treinar o modelo

Treine o modelo 2B usando um tamanho de lote de 16.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      cd transformers/
      export LD_LIBRARY_PATH=/usr/local/lib/
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      python3 examples/pytorch/xla_spawn.py \
         --num_cores=4 \
         examples/pytorch/language-modeling/run_clm.py \
         --num_train_epochs=3 \
         --dataset_name=wikitext \
         --dataset_config_name=wikitext-2-raw-v1 \
         --per_device_train_batch_size=16 \
         --per_device_eval_batch_size=16 \
         --do_train \
         --do_eval \
         --logging_dir=./tensorboard-metrics \
         --cache_dir=./cache_dir \
         --output_dir=/tmp/test-clm \
         --overwrite_output_dir \
         --cache_dir=/tmp \
         --config_name=examples/pytorch/language-modeling/my_config_2.json \
         --tokenizer_name=gpt2 \
         --block_size=1024 \
         --optim=adafactor \
         --adafactor=true \
         --save_strategy=no \
         --logging_strategy=no \
         --fsdp=full_shard \
         --fsdp_config=examples/pytorch/language-modeling/fsdp_config.json'

Excluir a TPU e o recurso na fila

Exclua a TPU e o recurso na fila no final da sessão.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Resultado do comparativo de mercado

O script de treinamento foi executado nas versões v5litepod-4, v5litepod-16 e v5litepod-64. A tabela a seguir mostra as capacidades de comparação para diferentes tipos de acelerador.

	v5litepod-4	v5litepod-16	v5litepod-64
Época	3	3	3
config	600 milhões	2 bilhões	16 bi
Tamanho global do lote	64	128	256
Capacidade de processamento (exemplos/segundo)	66	77	31

Treinar ViT na v5e

Este tutorial mostra como executar o VIT na v5e usando o repositório HuggingFace no PyTorch/XLA no conjunto de dados cifar10 (links em inglês).

Configurar

1. Crie variáveis de ambiente:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=tpu-name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d
   

2. Crie um recurso de TPU:

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   Observação: a sinalização QUOTA_TYPE pode ser reserved ou best-effort. Se nenhum deles for especificado, o padrão será on-demand. Acesse quotas para mais informações sobre os diferentes tipos de cotas compatíveis com o Cloud TPU.

   Será possível estabelecer uma conexão SSH com sua VM da TPU quando o QueuedResource estiver no estado ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   Quando o recurso na fila está no estado ACTIVE, a saída é semelhante a esta:

    state: ACTIVE
   

3. Instalar as dependências do PyTorch/XLA

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all
      --command='
         sudo apt-get update -y
         sudo apt-get install libomp5 -y
         pip3 install mkl mkl-include
         pip3 install tf-nightly tb-nightly tbp-nightly
         pip3 install numpy
         sudo apt-get install libopenblas-dev -y
         pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
         pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
   

4. Faça o download do repositório HuggingFace e instale os requisitos.

      gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command="
         git clone https://github.com/suexu1025/transformers.git vittransformers; \
         cd vittransformers; \
         pip3 install .; \
         pip3 install datasets; \
         wget https://github.com/pytorch/xla/blob/master/scripts/capture_profile.py"
   

Treinar o modelo

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export TF_CPP_MIN_LOG_LEVEL=0
      export XLA_USE_BF16=1
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      cd vittransformers
      python3 -u examples/pytorch/xla_spawn.py --num_cores 4 examples/pytorch/image-pretraining/run_mae.py --dataset_name=cifar10 \
      --remove_unused_columns=False \
      --label_names=pixel_values \
      --mask_ratio=0.75 \
      --norm_pix_loss=True \
      --do_train=true \
      --do_eval=true \
      --base_learning_rate=1.5e-4 \
      --lr_scheduler_type=cosine \
      --weight_decay=0.05 \
      --num_train_epochs=3 \
      --warmup_ratio=0.05 \
      --per_device_train_batch_size=8 \
      --per_device_eval_batch_size=8 \
      --logging_strategy=steps \
      --logging_steps=30 \
      --evaluation_strategy=epoch \
      --save_strategy=epoch \
      --load_best_model_at_end=True \
      --save_total_limit=3 \
      --seed=1337 \
      --output_dir=MAE \
      --overwrite_output_dir=true \
      --logging_dir=./tensorboard-metrics \
      --tpu_metrics_debug=true'

Excluir a TPU e o recurso na fila

Exclua a TPU e o recurso na fila no final da sessão.

gcloud compute tpus tpu-vm delete ${TPU_NAME}
   --project=${PROJECT_ID}
   --zone=${ZONE}
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
   --project=${PROJECT_ID}
   --zone=${ZONE}
   --quiet

Resultado do comparativo de mercado

A tabela a seguir mostra as capacidades de processamento do comparativo de mercado para diferentes tipos de acelerador.

	v5litepod-4	v5litepod-16	v5litepod-64
Época	3	3	3
Tamanho global do lote	32	128	512
Capacidade de processamento (exemplos/segundo)	201	657	2.844

TensorFlow 2.x

As seções a seguir descrevem exemplos de como treinar modelos do TensorFlow 2.x na TPU v5e.

Treinar Resnet em um host único v5e

Neste tutorial, descrevemos como treinar o ImageNet em v5litepod-4 ou v5litepod-8 usando um conjunto de dados fictício. Se você quiser usar um conjunto de dados diferente, consulte Como preparar o conjunto de dados.

Configurar

1. Crie variáveis de ambiente:

   export PROJECT_ID=your-project-ID
   export ACCELERATOR_TYPE=v5litepod-4
   export ZONE=us-east1-c
   export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pjrt
   export TPU_NAME=your-tpu-name
   export QUEUED_RESOURCE_ID=your-queued-resource-id
   export QUOTA_TYPE=quota-type
   

   ACCELERATOR_TYPE pode ser v5litepod-4 ou v5litepod-8.

2. Crie um recurso de TPU:

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --${QUOTA_TYPE}
   

   Observação: a sinalização QUOTA_TYPE pode ser reserved ou best-effort. Se nenhum deles for especificado, o padrão será on-demand. Acesse quotas para mais informações sobre os diferentes tipos de cotas compatíveis com o Cloud TPU.

   Será possível estabelecer uma conexão SSH com sua VM da TPU quando o recurso na fila estiver no estado ACTIVE. Para verificar o estado do recurso na fila, use o seguinte comando:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

3. Conectar-se à TPU usando SSH

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

4. Defina algumas variáveis de ambiente

   export MODELS_REPO=/usr/share/tpu/models
   export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
   export MODEL_DIR=gcp-directory-to-store-model
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
   export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so
   

5. Altere para o diretório do repositório de modelos e os requisitos de instalação.

   cd ${MODELS_REPO} && git checkout r2.15.0
   pip install -r official/requirements.txt
   

Treinar o modelo

Execute o script de treinamento.

python3 official/vision/train.py \
   --tpu=local \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

Excluir a TPU e o recurso na fila

1. Excluir a TPU

   gcloud compute tpus tpu-vm delete ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet
   

2. Excluir a solicitação de recurso na fila

   gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet
   

Treinar Resnet em um v5e de vários hosts

Neste tutorial, descrevemos como treinar o ImageNet em v5litepod-16 ou maior usando um conjunto de dados fictício. Se você quiser usar um conjunto de dados diferente, consulte Como preparar o conjunto de dados.

1. Crie variáveis de ambiente:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-east1-c
   export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pod-pjrt
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your-queued-resource-id
   export QUOTA_TYPE=quota-type
   

   ACCELERATOR_TYPE pode ser v5litepod-16 ou maior.

2. Crie um recurso de TPU:

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --${QUOTA_TYPE}
   

   Observação: a sinalização QUOTA_TYPE pode ser reserved ou best-effort. Se nenhum deles for especificado, o padrão será on-demand. Acesse quotas para mais informações sobre os diferentes tipos de cotas compatíveis com o Cloud TPU.

   Será possível estabelecer uma conexão SSH com sua VM da TPU quando o recurso na fila estiver no estado ACTIVE. Para verificar o estado do recurso na fila, use o seguinte comando:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

3. Conecte-se à TPU (worker zero) usando SSH.

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

4. Defina algumas variáveis de ambiente

   export MODELS_REPO=/usr/share/tpu/models
   export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
   export MODEL_DIR=gcp-directory-to-store-model
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export TPU_LOAD_LIBRARY=0
   export TPU_NAME=your_tpu_name
   

5. Altere para o diretório do repositório de modelos e os requisitos de instalação.

    cd $MODELS_REPO && git checkout r2.15.0
    pip install -r official/requirements.txt
   

Treinar o modelo

Execute o script de treinamento.

python3 official/vision/train.py \
   --tpu=${TPU_NAME} \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*"

Excluir a TPU e o recurso na fila

1. Excluir a TPU

   gcloud compute tpus tpu-vm delete ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet
   

2. Excluir a solicitação de recurso na fila

   gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet