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 tutorial, há componentes faturáveis do Google Cloud Platform, entre eles:

  • Compute Engine
  • Cloud TPU

Use a calculadora de preços para gerar uma estimativa de custos baseada na projeção de uso. Os novos usuários do GCP 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 Platform foi configurado corretamente.

  1. Faça login na sua Conta do Google.

    Se você ainda não tiver uma, inscreva-se.

  2. Selecione ou crie um projeto do Google Cloud Platform.

    Acessar a página Gerenciar recursos

  3. Verifique se o faturamento foi ativado no projeto do Google Cloud Platform.

    Saiba como ativar o faturamento

  4. Este tutorial inclui componentes faturáveis do Google Cloud Platform. Consulte a página de preços do 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

Acesse o console e crie uma nova instância de VM, fornecendo um nome e uma zona para a VM.

Se você planeja treinar o Resnet50 em dados reais, escolha o tipo de máquina com o maior número de CPUs possível. Normalmente, o Resnet50 é altamente vinculado à entrada, de modo que o treinamento pode ser muito lento, a menos que haja muitos workers para alimentar dados e RAM suficiente para manter um grande número de linhas de execução de worker. Para melhores resultados, selecione o tipo de máquina n1-highmem-96.

Selecione Debian GNU/Linux 9 Stretch + PyTorch/XLA para o disco de inicialização. Se você planeja fazer o download do ImageNet, especifique um tamanho de disco de pelo menos 300 GB. Se você planeja usar apenas dados falsos, é possível especificar o tamanho de disco padrão de 20 GB.

Neste tutorial, sugerimos o uso de ambos os conjuntos de dados.

Criar TPU

  1. Acesse o console para criar uma TPU.
  2. Em "Nome", especifique um nome para o pod de TPU.
  3. Em “Zona”, especifique a zona a ser usada para o Cloud TPU. Verifique se está na mesma zona que sua VM.
  4. Em "Tipo de TPU", selecione o tipo de Cloud TPU. Para melhores resultados, selecione a TPU v3-8.
  5. Em "Versão do software de TPU", selecione a versão estável mais recente (pytorch-0.5).
  6. Use a rede padrão.
  7. Defina o intervalo de endereços IP. Por exemplo, 10.240.0.0.

Treinar com o conjunto fake_data

Recomendamos o uso do conjunto de dados falso para sua execução inicial em vez do conjunto real do ImageNet, porque fake_data é instalado automaticamente na sua VM e requer menos tempo e menos recursos para o processamento.

Você pode usar o conda ou o Docker para o treinamento. Use o conda se não tiver familiaridade com o Docker.

  • Execute a versão do código empacotado em um ambiente conda dentro do disco Debian GNU/Linux 9 Stretch + PyTorch/XLA.
  • Execute usando o Docker. Levará alguns minutos para extrair a imagem mais recente do Docker antes de você começar o treinamento.

Opção 1

# Fill in your the name of your VM and the zone.
$ gcloud beta compute  ssh "your-VM-name" --zone "your-zone".
(vm)$ export TPU_IP_ADDRESS=your-ip-address
(vm)$ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
(vm)$ conda activate torch-xla-0.5
(torch-xla-0.5)$  python /usr/share/torch-xla-0.5/pytorch/xla/test/test_train_imagenet.py --fake_data --model=resnet50 --num_epochs=10 --num_workers=64 --batch_size=128 --log_steps=20

Opção 2

# Fill in your the name of your VM and the zone.
$ gcloud beta compute ssh "your-VM-name" --zone "your-zone".
(vm)$ export TPU_IP_ADDRESS=your-ip-address
(vm)$ docker run --shm-size 16G -e XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470" gcr.io/tpu-pytorch/xla:r0.5 python3 pytorch/xla/test/test_train_imagenet.py --model=resnet50 --num_epochs=10 --num_workers=64 --fake_data

Treinar com o conjunto de dados real

Se estiver tudo certo com o uso da sinalização --fake_data, você pode tentar usar o treinamento em dados reais, como do ImageNet.

Em geral, test_train_imagenet.py usa torchvision.datasets.ImageFolder para que você possa usar qualquer conjunto de dados estruturado adequadamente. Consulte a documentação de ImageFolder.

Algumas modificações sugeridas para dados reais, supondo que você tenha armazenado o conjunto de dados em ~/imagenet:

Opção 1

# Fill in your the name of your VM and the zone.
$ gcloud beta compute  ssh "your-VM-name" --zone "your-zone".
(vm)$ export TPU_IP_ADDRESS=your-ip-address
(vm)$ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
(vm)$ ulimit -n 10240
(vm)$ conda activate torch-xla-0.5
(torch-xla-0.5)$ python /usr/share/torch-xla-0.5/pytorch/xla/test/test_train_imagenet.py --datadir=~/imagenet --model=resnet50 --num_epochs=90 --num_workers=64 --batch_size=128 --log_steps=200

Opção 2

# Fill in your the name of your VM and the zone.
$ gcloud beta compute ssh "your-VM-name" --zone "your-zone".
(vm)$ export TPU_IP_ADDRESS=your-ip-address
(vm)$ docker run --shm-size 128G -v ~/imagenet:/tmp/imagenet -e XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470" gcr.io/tpu-pytorch/xla:r0.5 python3 pytorch/xla/test/test_train_imagenet.py --model=resnet50 --num_epochs=90 --num_workers=64 --log_steps=200 --datadir=/tmp/imagenet

No conjunto de dados completo do ImageNet com uma VM n1-highmem-96 e v3-8 TPU, o treinamento normalmente leva cerca de 20 minutos para a primeira época e cerca de 11 minutos para épocas subsequentes. O modelo deve alcançar aproximadamente 76% de precisão no primeiro nível em 90 épocas.

Solicitar cota adicional de CPU

Planeje e solicite mais recursos com pelo menos alguns dias de antecedência para garantir que haja tempo suficiente para atender ao seu pedido.

  1. Acesse a página Cotas.

    Acessar a página "Cotas"

  2. No menu Serviço, selecione API Cloud TPU.
  3. Selecione a região ou as zonas em que você quer usar as CPUs.
  4. No menu Métrica, selecione "Nenhuma" e insira "CPUs" na caixa de pesquisa.
  5. Selecione CPUs.
  6. Na lista, selecione API Compute Engine API - CPUs e clique em Editar cotas na parte superior da página.
  7. Insira o valor da cota que você está solicitando e uma descrição (obrigatória). Em seguida, clique em Concluído. Uma solicitação é enviada ao seu provedor de serviços para aprovação.

Limpeza

Para evitar cobranças dos recursos usados neste tutorial na conta do Google Cloud Platform:

  1. Encerre a conexão com a VM do Compute Engine.
  2. Exclua a VM.
  3. Exclua a TPU.

A seguir

Use as colabs do PyTorch:

Esta página foi útil? Conte sua opinião sobre:

Enviar comentários sobre…

Precisa de ajuda? Acesse nossa página de suporte.