Entrena Resnet50 en Cloud TPU con PyTorch

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

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

Advertencia: En este instructivo, se usa un conjunto de datos de terceros. Google no hace ninguna declaración ni proporciona garantía alguna acerca de la validez o cualquier otro aspecto de este conjunto de datos.

Objetivos

• Preparar el conjunto de datos
• Ejecutar el trabajo de entrenamiento
• Verificar los resultados de salida

Costos

En este documento, usarás los siguientes componentes facturables de Google Cloud:

• Compute Engine
• Cloud TPU

Para generar una estimación de costos en función del uso previsto, usa la calculadora de precios.

Es posible que los usuarios nuevos de Google Cloud cumplan con los requisitos para acceder a una prueba gratuita.

Antes de comenzar

Antes de comenzar este instructivo, verifica que tu proyecto de Google Cloud esté configurado de forma correcta.

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

En esta explicación, se usan componentes facturables de Google Cloud. Consulta la página de precios de Cloud TPU para calcular los costos. Asegúrate de realizar una limpieza de los recursos que crees cuando hayas terminado de usarlos para evitar cargos innecesarios.

Crea una VM de TPU

1. Abre una ventana de Cloud Shell.

   Abrir Cloud Shell

2. Crea 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 ejecutes un comando en una VM de Cloud Shell nueva, se mostrará la página Authorize Cloud Shell. Haz clic en Authorize en la parte inferior de la página para permitir que gcloud realice llamadas a la API de Google Cloud con tus credenciales.

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

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

4. Instala PyTorch/XLA en la VM de tu 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 de 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

Realiza una limpieza

Para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos usados en este instructivo, borra el proyecto que contiene los recursos, o consérvalo y borra los recursos individuales.

1. Desconéctate de la VM de TPU:

   (vm) $ exit

   La instrucción ahora debería mostrar username@projectname, que indica que estás en Cloud Shell.

2. Borra la VM de tu TPU.

   $ gcloud compute tpus tpu-vm delete your-tpu-name \
      --zone=us-central1-a

¿Qué sigue?

• Entrena modelos de difusión con PyTorch
• Soluciona problemas de PyTorch en TPU
• Documentación de PyTorch/XLA