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ê 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 o tutorial, verifique se o projeto do Google Cloud foi configurado corretamente.

  1. Faça login na sua conta do Google Cloud. Se você começou a usar o Google Cloud agora, crie uma conta para avaliar o desempenho de nossos produtos em situações reais. Clientes novos também recebem US$ 300 em créditos para executar, testar e implantar cargas de trabalho.

  2. No console do Google Cloud, na página do seletor de projetos, selecione ou crie um projeto do Google Cloud.

    Acessar o seletor de projetos

  3. Verifique se a cobrança está ativada para o seu projeto do Google Cloud. Saiba como verificar se o faturamento está ativado em um projeto.

  4. No console do Google Cloud, na página do seletor de projetos, selecione ou crie um projeto do Google Cloud.

    Acessar o seletor de projetos

  5. Verifique se a cobrança está ativada para o seu projeto do Google Cloud. Saiba como verificar se o faturamento está ativado em um projeto.

    6. Este tutorial usa 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 apagar os recursos criados ao terminar de usá-los.

Criar uma VM da TPU

  1. Abra uma janela do Cloud Shell.

    Abrir o Cloud Shell

  2. Criar uma VM da TPU 

    gcloud compute tpus tpu-vm create your-tpu-name \
--accelerator-type=v4-8 \
--version=tpu-vm-v4-pt-2.0 \
--zone=us-central2-b \
--project=your-project

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

    gcloud compute tpus tpu-vm ssh  your-tpu-name --zone=us-central2-b

  4. Clone o repositório PyTorch/XLA no GitHub

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

  5. Executar o script de treinamento com dados falsos

    Há duas opções de ambiente de execução do PyTorch/XLA: PJRT e XRT. Recomendamos usar PJRT, a menos que você saiba que precisa usar XRT. Para saber mais sobre as diferentes configurações do ambiente de execução, consulte a documentação do ambiente de execução do PJRT.

    PJRT

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

    XRT (legado)

    Configurar XRT 

     (vm) $ export TPU_NUM_DEVICES=4
 (vm) $ export XRT_TPU_CONFIG="localservice;0;localhost:51011"

    Executar o script de treinamento

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

Se você conseguir treinar o modelo usando dados falsos, tente treinar em dados reais, como o ImageNet. Para ver instruções sobre como fazer o download do ImageNet, consulte Como fazer o download do ImageNet. No comando do script de treinamento, a sinalização --datadir especifica o local do conjunto de dados em que o treinamento será feito. O comando a seguir pressupõe que o conjunto de dados do ImageNet está localizado em ~/imagenet.

PJRT

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

XRT (legado)

Configurar XRT 

 <pre class="lang-sh prettyprint">
 <span class="no-select">(vm) $ </span>export TPU_NUM_DEVICES=4
 <span class="no-select">(vm) $ </span>export XRT_TPU_CONFIG="localservice;0;localhost:51011"
 </pre>

Executar o script de treinamento

 <pre class="lang-sh prettyprint">
 <span class="no-select">(vm) $ </span>python3 xla/test/test_train_mp_imagenet.py --datadir=<var>~/imagenet</var> --batch_size=256 --num_epochs=1
 </pre>

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. Desconecte-se da 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 resnet50-tutorial \
   --zone=us-central2-b

A seguir

Teste as colabs do PyTorch: