Como treinar o Resnet50 no Cloud TPU com PyTorch

Neste tutorial, mostramos como treinar o modelo ResNet-50 em um dispositivo Cloud TPU com PyTorch. É possível aplicar o mesmo padrão a outros modelos de classificação de imagem otimizados para TPU que usam o PyTorch e o conjunto de dados do ImageNet.

O modelo deste tutorial é baseado no artigo Deep Residual Learning for Image Recognition, que foi o primeiro a apresentar a arquitetura de rede residual ou ResNet. O tutorial usa a variante de 50 camadas, ResNet-50, e demonstra o treinamento do modelo usando PyTorch/XLA.

Objetivos

  • Preparar o conjunto de dados.
  • Executar o job de treinamento.
  • Verificar os resultados da saída.

Custos

Neste documento, você vai usar os seguintes componentes faturáveis do Google Cloud:

  • Compute Engine
  • Cloud TPU

Para gerar uma estimativa de custo baseada na projeção de uso deste tutorial, use a calculadora de preços.

Novos usuários do Google Cloud podem estar qualificados para uma avaliação gratuita.

Antes de começar

Antes de começar este tutorial, verifique se o projeto do Google Cloud foi configurado corretamente.

  1. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

  2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  3. Make sure that billing is enabled for your Google Cloud project.

  4. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  5. Make sure that billing is enabled for your Google Cloud project.

    6. Neste tutorial, usamos componentes faturáveis do Google Cloud. Consulte a página de preços da Cloud TPU para fazer uma estimativa dos custos. Para evitar cobranças desnecessárias, não se esqueça de limpar os recursos criados ao terminar de usá-los.

Criar uma VM de TPU

  1. Abra uma janela do Cloud Shell.

    Abra o Cloud Shell

  2. Criar uma VM de TPU

    gcloud compute tpus tpu-vm create your-tpu-name \
--accelerator-type=v3-8 \
--version=tpu-ubuntu2204-base \
--zone=us-central1-a \
--project=your-project

  3. Conecte-se à VM da TPU usando SSH:

    gcloud compute tpus tpu-vm ssh  your-tpu-name --zone=us-central1-a

  4. Instale o PyTorch/XLA na VM de TPU:

    (vm)$ pip install torch torch_xla[tpu] torchvision -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html

  5. Clone o repositório do PyTorch/XLA no GitHub.

    (vm)$ git clone --depth=1 https://github.com/pytorch/xla.git

  6. Executar o script de treinamento com dados falsos

    (vm) $ PJRT_DEVICE=TPU python3 xla/test/test_train_mp_imagenet.py --fake_data --batch_size=256 --num_epochs=1

Limpar

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados no tutorial, exclua o projeto que os contém ou mantenha o projeto e exclua os recursos individuais.

  1. Encerre a conexão com a VM da TPU:

    (vm) $ exit

    Agora, o prompt precisa ser username@projectname, mostrando que você está no Cloud Shell.

  2. Exclua a VM da TPU.

    $ gcloud compute tpus tpu-vm delete your-tpu-name \
   --zone=us-central1-a

A seguir