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.

Alerta: neste tutorial, é usado um conjunto de dados de terceiros. O Google não oferece declarações, proteções ou outras garantias sobre a validade ou quaisquer outros aspectos desse conjunto de dados.

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, entre eles:

• Compute Engine
• Cloud TPU

Use a calculadora de preços para gerar uma estimativa de custos baseada na projeção de uso. 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.

   Observação: se você não pretende manter os recursos criados neste procedimento, crie um projeto novo em vez de selecionar um que já existe. Depois de concluir essas etapas, é possível excluir o projeto. Para fazer isso, basta remover todos os recursos associados a ele.

   Acessar o seletor de projetos

3. Verifique se o faturamento está ativado para seu projeto na nuvem. Saiba como confirmar se o faturamento está ativado para o projeto.

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.

Configurar uma instância do Compute Engine

1. Abra uma janela do Cloud Shell.

   Abrir o Cloud Shell

2. Crie uma variável para o ID do seu projeto.

   export PROJECT_ID=project-id
   

3. Configure a ferramenta de linha de comando da gcloud para usar o projeto em que a Cloud TPU será criada.

   gcloud config set project ${PROJECT_ID}
   

   Na primeira vez que você executar esse comando em uma nova VM do Cloud Shell, será exibida uma página Authorize Cloud Shell. Clique em Authorize na parte inferior da página para permitir que gcloud faça chamadas de API do GCP com suas credenciais.

4. Definir a versão do PyTorch

   export VERSION=1.8
   

5. Inicie o recurso do Compute Engine exigido para este tutorial.

   gcloud compute instances create resnet50-tutorial \
   --zone=us-central1-a \
   --machine-type=n1-highmem-96  \
   --image-family=torch-xla \
   --image-project=ml-images  \
   --boot-disk-size=300GB \
   --scopes=https://www.googleapis.com/auth/cloud-platform
   

   Observação: ao executar este comando, você verá um alerta:

   Disk size: '300 GB' is larger than image size
   

   É possível ignorar a mensagem de alerta com segurança.

   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.

   Observação: verifique sua cota de CPI para garantir que você tem um número suficiente de CPUs disponíveis. Para aumentar sua cota, consulte cota de CPU.

   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.

Iniciar um recurso da Cloud TPU

1. Inicie um recurso do Cloud TPU na máquina virtual do Compute Engine, usando o seguinte comando:

   (vm) $ gcloud compute tpus create resnet50-tutorial \
   --zone=us-central1-a \
   --network=default \
   --version=pytorch-${VERSION?}  \
   --accelerator-type=v3-8
   

2. Identifique o endereço IP do recurso do Cloud TPU.

   (vm) $ gcloud compute tpus list --zone=us-central1-a
   

   Importante: o endereço IP está localizado na coluna NETWORK_ENDPOINTS. Você precisará desse endereço IP, sem o número da porta, ao criar e configurar o ambiente PyTorch.

Crie e configure o ambiente PyTorch

1. Conecte-se à nova instância do Compute Engine.

   gcloud compute ssh resnet50-tutorial --zone=us-central1-a
   

Daqui em diante, o prefixo (vm)$ significa que é preciso executar o comando na instância de VM do Compute Engine.

1. Inicie um ambiente conda.

   (vm) $ export VERSION=1.8
   (vm) $ conda activate torch-xla-${VERSION?}
   

2. Configure variáveis de ambiente para o recurso Cloud TPU.

   Observação: a variável TPU_IP_ADDRESS precisa ser igual ao endereço IP da Cloud TPU identificado quando você lançou o recurso Cloud TPU sem o número da porta. É o endereço IP interno exibido no Console do Google Cloud, em Compute Engine → TPUs.

   (vm) $ export TPU_IP_ADDRESS=ip-address
   

   (vm) $ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
   

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.

(vm) $ python /usr/share/torch-xla-${VERSION?}/pytorch/xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=2 --batch_size=128 --log_steps=20

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_mp_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 de linha de comando sugeridas para o uso de dados reais, supondo que você tenha armazenado o conjunto de dados em ~/imagenet:

(vm) $ python /usr/share/torch-xla-${VERSION?}/pytorch/xla/test/test_train_mp_imagenet.py --datadir=~/imagenet --model=resnet50 --num_epochs=90 --num_workers=8 --batch_size=128 --log_steps=200

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 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 instância do Compute Engine, caso ainda não tenha feito isso:

   (vm)$ exit
   

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

2. No Cloud Shell, execute ctpu delete com as sinalizações --zone e --name usadas ao configurar a VM do Compute Engine e a Cloud TPU. Isso exclui sua VM e sua Cloud TPU.

   $ ctpu delete --project=${PROJECT_ID} \
     --name=resnet50-tutorial \
     --zone=us-central1-a
   

   Importante: se você definir o nome dos recursos de TPU ao executar ctpu up, especifique esse nome com a sinalização --name ao executar ctpu delete para encerrar seus recursos de TPU.

3. Execute ctpu status para garantir que não haja instâncias alocadas e evitar cobranças desnecessárias no uso da TPU. A exclusão pode levar vários minutos. Uma resposta como esta indica que não há mais instâncias alocadas:

   $ ctpu status --project=${PROJECT_ID} \
     --zone=us-central1-a
   

   2018/04/28 16:16:23 WARNING: Setting zone to "us-central1-a"
   No instances currently exist.
           Compute Engine VM:     --
           Cloud TPU:             --
   

Observação: para saber mais sobre limites de armazenamento gratuito e outras informações sobre preços, consulte o guia de preços do Cloud Storage.

A seguir

Teste as colabs do PyTorch:

• Primeiros passos com o PyTorch em Cloud TPUs
• Como treinar MNIST em TPUs (em inglês)
• Como treinar o ResNet18 em TPUs com o conjunto de dados Cifar10 (em inglês)
• Inferência com modelo ResNet50 pré-treinado (em inglês)
• Transferência rápida de estilo neural
• Treinamento de vários núcleos da AlexNet no Fashion MNIST
• Treinamento de núcleo único da AlexNet no Fashion MNIST