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.

    Si todavía no tienes una cuenta, regístrate para obtener una nueva.

  2. En la página de selección de proyectos de Cloud Console, selecciona o crea un proyecto de Cloud.

    Ir a la página Selector de proyectos

  3. Asegúrate de que la facturación esté habilitada para tu proyecto de Google Cloud. Obtén información sobre cómo confirmar que tienes habilitada la facturación para tu 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.

    Abre 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 de gcloud para usar el proyecto en el que deseas crear Cloud TPU.

    gcloud config set project ${PROJECT_ID}
    
  4. Configura la versión de PyTorch

    export VERSION=1.6
    
  5. 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-${VERSION?}  \
    --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) $ export VERSION=1.6
    (vm) $ conda activate torch-xla-${VERSION?}
    
  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-${VERSION?}/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 línea de comandos para usar datos reales, bajo la suposición de que almacenaste el conjunto de datos en ~/imagenet:

(vm) $ python /usr/share/torch-xla-${VERSION?}/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.

Limpia

Sigue estos pasos para evitar que se apliquen cargos a tu cuenta de Google Cloud Platform por los recursos que usaste en este instructivo:

  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 --name=resnet50-tutorial \
      --zone=europe-west4-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 --zone=europe-west4-a
    
    2018/04/28 16:16:23 WARNING: Setting zone to "europe-west4-a"
    No instances currently exist.
            Compute Engine VM:     --
            Cloud TPU:             --
    

Qué sigue

Prueba los siguientes colaboradores de PyTorch: