Formação da Resnet50 na Cloud TPU com o PyTorch

Este tutorial mostra como preparar o modelo ResNet-50 num dispositivo Cloud TPU com o PyTorch. Pode aplicar o mesmo padrão a outros modelos de classificação de imagens otimizados para TPUs que usam o PyTorch e o conjunto de dados ImageNet.

O modelo neste tutorial baseia-se na aprendizagem residual profunda para reconhecimento de imagens, que introduz primeiro a arquitetura de rede residual (ResNet). O tutorial usa a variante de 50 camadas, o ResNet-50, e demonstra a preparação do modelo usando o PyTorch/XLA.

Crie uma VM da TPU

  1. Abra uma janela do Cloud Shell.

    Abra o Cloud Shell

  2. Crie uma VM da TPU

    gcloud compute tpus tpu-vm create your-tpu-name \
--accelerator-type=v3-8 \
--version=tpu-ubuntu2204-base \
--zone=us-central1-a \
--project=your-project

  3. Estabeleça ligação à VM de TPU através do SSH:

    gcloud compute tpus tpu-vm ssh  your-tpu-name --zone=us-central1-a

  4. Instale o PyTorch/XLA na sua VM da TPU:

    (vm)$ pip install torch torch_xla[tpu] torchvision -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html

  5. Clone o repositório do GitHub PyTorch/XLA

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

  6. Execute o script de preparação com dados falsos

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