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.
-
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.
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-ubuntu2204-base \ --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
Instale o PyTorch/XLA na VM da TPU:
(vm)$ pip install torch~=2.3.0 torch_xla[tpu]~=2.3.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html
Clone o repositório do GitHub PyTorch/XLA
(vm)$ git clone --depth=1 --branch r2.3 https://github.com/pytorch/xla.git
Executar o script de treinamento com dados fictícios
(vm) $ PJRT_DEVICE=TPU python3 xla/test/test_train_mp_imagenet.py --fake_data --batch_size=256 --num_epochs=1
Se você conseguir treinar o modelo usando dados falsos, tente treinar
dados reais, como o ImageNet. Para instruções sobre como fazer o download do ImageNet, consulte
Como fazer o download do ImageNet. No comando do script de treinamento,
a flag --datadir
especifica o local do conjunto de dados no qual treinar.
O comando a seguir pressupõe que o conjunto de dados do ImageNet está localizado em ~/imagenet
.
(vm) $ PJRT_DEVICE=TPU python3 xla/test/test_train_mp_imagenet.py --datadir=~/imagenet --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.
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