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.
Antes de começar
Antes de começar o tutorial, verifique se o projeto do Google Cloud foi configurado corretamente.
- 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.
-
No console do Google Cloud, na página do seletor de projetos, selecione ou crie um projeto do Google Cloud.
-
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.
-
No console do Google Cloud, na página do seletor de projetos, selecione ou crie um projeto do Google Cloud.
-
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.
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
Abra uma janela do Cloud Shell.
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
Conecte-se à VM da TPU usando SSH:
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=us-central2-b
Clone o repositório PyTorch/XLA no GitHub
(vm)$ git clone --depth=1 --branch r2.0 https://github.com/pytorch/xla.git
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.
Desconecte-se da VM da TPU:
(vm) $ exit
Agora, o prompt precisa ser
username@projectname
, mostrando que você está no Cloud Shell.Exclua a VM da TPU.
$ gcloud compute tpus tpu-vm delete resnet50-tutorial \ --zone=us-central2-b
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