Entrenar Resnet50 en TPU de Cloud con PyTorch

En este tutorial se explica cómo entrenar el modelo ResNet-50 en un dispositivo TPU de Cloud con PyTorch. Puedes aplicar el mismo patrón a otros modelos de clasificación de imágenes optimizados para TPUs que usen PyTorch y el conjunto de datos ImageNet.

El modelo de este tutorial se basa en Aprendizaje residual profundo para el reconocimiento de imágenes, que introduce por primera vez la arquitectura de la red residual (ResNet). En el tutorial se usa la variante de 50 capas, ResNet-50, y se muestra cómo entrenar el modelo con PyTorch/XLA.

Advertencia: En este tutorial se utiliza un conjunto de datos de terceros. Google no proporciona representación ni garantías de ningún tipo en relación con la validez ni con ningún otro aspecto de este conjunto de datos.

Crear una VM de TPU

1. Abre una ventana de Cloud Shell.

   Abrir Cloud Shell

2. Crear una VM de 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

   Nota: La primera vez que ejecutas un comando en una nueva VM de Cloud Shell, se muestra una página de Authorize Cloud Shell. Haz clic en Authorize en la parte inferior de la página para permitir que gcloud haga llamadas a la API con tus credenciales.
   Google Cloud

3. Conéctate a tu VM de TPU mediante SSH:

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

4. Instala PyTorch/XLA en tu VM de 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. Clona el repositorio de GitHub PyTorch/XLA.

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

6. Ejecuta la secuencia de comandos de entrenamiento con datos falsos

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