Entrena Resnet50 en Cloud TPU con PyTorch

En este instructivo, se muestra cómo entrenar el modelo ResNet-50 en un dispositivo 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 ImageNet.

El modelo de este instructivo se basa en Aprendizaje residual profundo para el reconocimiento de imágenes, que 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.

Objetivos

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

Costos

En este instructivo, se usan componentes facturables de Google Cloud, que incluyen los siguientes:

  • Compute Engine
  • Cloud TPU

Usa la calculadora de precios para generar una estimación de los costos según el uso previsto. Los usuarios nuevos de Google Cloud pueden ser elegibles para obtener una prueba gratuita.

Antes de comenzar

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

  1. Accede a tu cuenta de Google Cloud. Si eres nuevo en Google Cloud, crea una cuenta para evaluar el rendimiento de nuestros productos en situaciones reales. Los clientes nuevos también obtienen $300 en créditos gratuitos para ejecutar, probar y, además, implementar cargas de trabajo.
  2. En la página del selector de proyectos de Google Cloud Console, selecciona o crea un proyecto de Google Cloud.

    Ir al selector de proyecto

  3. Asegúrate de que la facturación esté habilitada para tu proyecto de Cloud. Descubre cómo confirmar que tienes habilitada la facturación en un proyecto.

  4. 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 limpiar los recursos que crees cuando hayas terminado de usarlos para evitar cargos innecesarios.

Configura una instancia de Compute Engine

  1. Abre una ventana de Cloud Shell.

    Abrir Cloud Shell

  2. Crea una variable para el ID de tu proyecto.

    export PROJECT_ID=project-id
    
  3. Configura la herramienta de línea de comandos gcloud para usar el proyecto en el que deseas crear Cloud TPU.

    gcloud config set project ${PROJECT_ID}
    

    La primera vez que ejecutes este 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 GCP con tus credenciales.

  4. Inicia el recurso de Compute Engine necesario para este instructivo.

    gcloud compute instances create resnet50-tutorial \
    --zone=us-central1-a \
    --machine-type=n1-highmem-96  \
    --image-family=torch-xla \
    --image-project=ml-images  \
    --boot-disk-size=300GB \
    --scopes=https://www.googleapis.com/auth/cloud-platform
    

    Si planeas entrenar Resnet50 con datos reales, elige el tipo de máquina con la mayor cantidad de CPU posible. Por lo general, Resnet50 depende en gran medida de las entradas, por lo que el entrenamiento puede ser bastante lento, a menos que haya muchos trabajadores que ingresen datos y suficiente RAM para mantener una gran cantidad de subprocesos de trabajador. Para obtener mejores resultados, selecciona el tipo de máquina n1-highmem-96.

    Si planeas descargar ImageNet, especifica un tamaño de disco de al menos 300 GB. Si planeas usar solo datos falsos, puedes especificar el tamaño de disco predeterminado en 20 GB. En este instructivo, se sugiere usar ambos conjuntos de datos.

Inicia un recurso de Cloud TPU

  1. En la máquina virtual de Compute Engine, inicia un recurso de Cloud TPU con el siguiente comando:

    (vm) $ gcloud compute tpus create resnet50-tutorial \
    --zone=us-central1-a \
    --network=default \
    --version=pytorch-1.9 \
    --accelerator-type=v3-8
    
  2. Identifica la dirección IP para el recurso Cloud TPU.

    (vm) $ gcloud compute tpus list --zone=us-central1-a
    

Crea y configura el entorno PyTorch

  1. Conéctate a la instancia nueva de Compute Engine.

    gcloud compute ssh resnet50-tutorial --zone=us-central1-a
    

A partir de este momento, el prefijo (vm)$ indica que debes ejecutar el comando en la instancia de VM de Compute Engine.

  1. Inicia un entorno conda.

    (vm) $ conda activate torch-xla-1.9
    
  2. Configura variables de entorno para el recurso de Cloud TPU.

    (vm) $ export TPU_IP_ADDRESS=ip-address
    
    (vm) $ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
    

Entrenamiento con el conjunto fake_data

Recomendamos usar el conjunto de datos falsos para la ejecución inicial, en lugar del conjunto real de ImageNet, ya que fake_data se instala de forma automática en la VM y requiere menos tiempo y procesamiento de menor cantidad de recursos.

(vm) $ python /usr/share/torch-xla-1.9/pytorch/xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=2 --batch_size=128 --log_steps=20

Entrenamiento con el conjunto de datos reales

Si todo funciona bien con la marca --fake_data, puedes intentar entrenar con datos reales, como ImageNet.

En general, test_train_mp_imagenet.py usa torchvision.datasets.ImageFolder para que puedas usar cualquier conjunto de datos que esté estructurado de forma correcta. Consulta la documentación de ImageFolder.

A continuación, se sugieren algunas modificaciones de la línea de comandos para usar datos reales, si es que almacenaste el conjunto de datos en ~/imagenet:

(vm) $ python /usr/share/torch-xla-1.9/pytorch/xla/test/test_train_mp_imagenet.py --datadir=~/imagenet --model=resnet50 --num_epochs=90 --num_workers=8 --batch_size=128 --log_steps=200

Solicita una cuota de CPU adicional

Planifica el uso de recursos y solicita recursos adicionales con varios días de antelación a fin de asegurarte de que haya tiempo suficiente para que se concrete tu pedido.

  1. Ve a la página Cuotas.

    Ir a la página Cuotas

  2. En el menú Servicio, selecciona API de Cloud TPU.
  3. Selecciona la región o zonas donde quieres usar las CPU.
  4. En el menú Métrica, selecciona Ninguno y luego ingresa CPU en el cuadro de búsqueda.
  5. Selecciona CPU.
  6. En la lista, selecciona API de Compute Engine - CPU y haz clic en Editar cuotas, en la parte superior de la página.
  7. Ingresa la cantidad de cuota que estás solicitando y una descripción (obligatorio), luego haz clic en Listo. Se envía una solicitud a tu proveedor de servicios para que la apruebe.

Realice 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 conserva el proyecto y borra los recursos individuales.

  1. Desconéctate de la instancia de Compute Engine, si aún no lo hiciste:

    (vm)$ exit
    

    El mensaje ahora debería mostrar user@projectname, que indica que estás en Cloud Shell.

  2. En Cloud Shell, ejecuta ctpu delete con las marcas --zone y --name que usaste cuando configuraste la VM de Compute Engine y Cloud TPU. Esto borra tanto tu VM como tu Cloud TPU.

    $ ctpu delete --project=${PROJECT_ID} \
      --name=resnet50-tutorial \
      --zone=us-central1-a
    
  3. Ejecuta ctpu status para asegurarte de no tener instancias asignadas y así evitar cargos innecesarios por el uso de TPU. La eliminación puede tomar varios minutos. Una respuesta como la que se muestra a continuación indica que no hay más instancias asignadas:

    $ ctpu status --project=${PROJECT_ID} \
      --zone=us-central1-a
    
    2018/04/28 16:16:23 WARNING: Setting zone to "us-central1-a"
    No instances currently exist.
            Compute Engine VM:     --
            Cloud TPU:             --
    

Qué sigue

Prueba los siguientes colaboradores de PyTorch: