Usa las TPU para entrenar tu modelo

Las unidades de procesamiento tensorial (TPU) son los ASIC de Google desarrollados de manera personalizada que se utilizan para acelerar las cargas de trabajo de aprendizaje automático. Puedes ejecutar tus trabajos de entrenamiento en AI Platform con Cloud TPU. AI Platform ofrece una interfaz de administración de trabajos para que no tengas que administrar la TPU tú mismo. En su lugar, puedes usar la API jobs de AI Platform de la misma forma en que la usas para el entrenamiento en una CPU o una GPU.

Las API de TensorFlow de nivel alto te ayudan a ejecutar los modelos en el hardware de Cloud TPU.

Configura y prueba el entorno de GCP

Configura el entorno de GCP a través de la sección de configuración de la guía de introducción.

Autoriza a Cloud TPU para que tenga acceso al proyecto

Sigue estos pasos para autorizar el nombre de la cuenta de servicio de Cloud TPU asociado con el proyecto de GCP:

  1. Llama a projects.getConfig para obtener el nombre de la cuenta de servicio de Cloud TPU. Ejemplo:

    curl -H "Authorization: Bearer $(gcloud auth print-access-token)"  \
        https://ml.googleapis.com/v1/projects/<your-project-id>:getConfig
    
  2. Guarda el valor del campo tpuServiceAccount que muestra la API.

A continuación, agrega la cuenta de servicio de Cloud TPU como miembro del proyecto con la función de Agente de servicios de Cloud ML. Completa los pasos siguientes en Google Cloud Platform Console o utiliza el comando gcloud:

Console

  1. Accede a Google Cloud Platform Console y elige el proyecto en el cual estés usando la TPU.
  2. Elige IAM y admin > IAM.
  3. Haz clic en el botón Agregar para agregar un miembro al proyecto.
  4. Ingresa la cuenta de servicio de la TPU en el cuadro de texto Miembros.
  5. Haz clic en la lista desplegable de Funciones.
  6. Habilita la función Agente de servicios de Cloud ML (Administración de servicios > Agente de servicios de Cloud ML).

gcloud

  1. Configura las variables de entorno que contengan el ID del proyecto y la cuenta de servicio de Cloud TPU:

    PROJECT_ID=your-project-id
    SVC_ACCOUNT=your-tpu-sa-123@your-tpu-sa.google.com.iam.gserviceaccount.com
    
  2. Otorga la función ml.serviceAgent a la cuenta de servicio de Cloud TPU:

    gcloud projects add-iam-policy-binding $PROJECT_ID \
        --member serviceAccount:$SVC_ACCOUNT --role roles/ml.serviceAgent
    

Para obtener más detalles sobre cómo otorgar funciones a las cuentas de servicios, consulta la documentación de Cloud IAM.

Ejecuta el modelo ResNet-50 de muestra

Por intermedio de esta sección, se te muestra cómo entrenar la referencia modelo ResNet-50 de TensorFlow mediante un conjunto de datos falsos que se proporciona en gs://cloud-tpu-test-datasets/fake_imagenet. El trabajo de ejemplo usa el nivel de escala BASIC_TPU predefinido para la configuración de la máquina. En próximas secciones de la guía, se muestra cómo definir una configuración personalizada.

Ejecuta los siguientes comandos para obtener el código y enviar el trabajo de entrenamiento a AI Platform:

  1. Descarga el código del modelo de referencia:

    mkdir tpu-demos && cd tpu-demos
    wget https://github.com/tensorflow/tpu/archive/r1.14.tar.gz
    tar -xzvf r1.14.tar.gz && rm r1.14.tar.gz
    
  2. Ve al directorio official dentro de la estructura de directorios descomprimidos:

    cd tpu-r1.14/models/official/
    
  3. Edita ./resnet/resnet_main.py y cambia el código para usar importaciones relativas explícitas cuando importas submódulos. Por ejemplo, cambia este código:

    import resnet_model
    

    por este código:

    from . import resnet_model
    

    Usa el patrón anterior para cambiar todas las importaciones en el archivo, y luego guarda el archivo.

  4. Busca cualquier importación de submódulo en otros archivos en la muestra y actualízalos para que también puedan usar las importaciones relativas explícitas.

  5. Configura algunas variables de entorno:

    JOB_NAME=tpu_1
    STAGING_BUCKET=gs://my_bucket_for_staging
    REGION=us-central1
    DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
    OUTPUT_PATH=gs://my_bucket_for_model_output
    

    Las siguientes regiones actualmente proporcionan acceso a las TPU:

    • us-central1

  6. Envía tu trabajo de entrenamiento con el comando gcloud ai-platform jobs submit training:

    gcloud ai-platform jobs submit training $JOB_NAME \
            --staging-bucket $STAGING_BUCKET \
            --runtime-version 1.14 \
            --scale-tier BASIC_TPU \
            --module-name resnet.resnet_main \
            --package-path resnet/ \
            --region $REGION \
            -- \
            --data_dir=$DATA_DIR \
            --model_dir=$OUTPUT_PATH
    

Más información sobre cómo entrenar un modelo en Cloud TPU

En la primera parte de esta guía, se muestra cómo usar el código de muestra de ResNet-50. En esta sección, se brinda más información sobre cómo configurar un trabajo y entrenar un modelo en AI Platform con Cloud TPU.

Especifica una región en la que se ofrezcan las TPU

Es necesario que ejecutes tu trabajo en una región en la que las TPU se encuentren disponibles. Las siguientes regiones actualmente proporcionan acceso a las TPU:

  • us-central1

A fin de comprender totalmente las regiones disponibles para los servicios de AI Platform, incluidos el entrenamiento de modelos y la predicción en línea o por lotes, consulta la guía de regiones.

Control de versiones de TensorFlow y AI Platform

Las versiones 1.12 , 1.13 y 1.14 del entorno de ejecución de AI Platform se encuentran disponibles para entrenar los modelos en Cloud TPU. Consulta más información sobre las versiones del entorno de ejecución de AI Platform y las versiones correspondientes de TensorFlow.

La política de control de versiones es la misma que para Cloud TPU. En la solicitud de trabajo de entrenamiento, asegúrate de especificar una versión de entorno de ejecución que esté disponible para las TPU y coincida con la versión de TensorFlow que se usó en el código de entrenamiento.

Conéctate con el servidor gRPC de la TPU

En el programa de TensorFlow, debes usar TPUClusterResolver para conectarte con el servidor gRPC de la TPU que se ejecuta en la VM de TPU. TPUClusterResolver muestra la dirección IP y el puerto de Cloud TPU.

El siguiente ejemplo muestra la manera en que el código de muestra de ResNet-50 usa TPUClusterResolver:

tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
    FLAGS.tpu,
    zone=FLAGS.tpu_zone,
    project=FLAGS.gcp_project)

config = tpu_config.RunConfig(
    cluster=tpu_cluster_resolver,
    model_dir=FLAGS.model_dir,
    save_checkpoints_steps=max(600, FLAGS.iterations_per_loop),
    tpu_config=tpu_config.TPUConfig(
        iterations_per_loop=FLAGS.iterations_per_loop,
        num_shards=FLAGS.num_cores,
        per_host_input_for_training=tpu_config.InputPipelineConfig.PER_HOST_V2))  # pylint: disable=line-too-long
.

Asigna operaciones en las TPU

Para usar las TPU en una máquina, debes utilizar la API de TPUEStimator de TensorFlow, que se hereda de la API de Estimator de TensorFlow de nivel alto.

  • TPUEstimator maneja muchos de los detalles de la ejecución en dispositivos TPU, como la replicación de entradas y modelos en relación con cada núcleo y el regreso al host periódicamente para ejecutar enlaces.
  • La API de TensorFlow de alto nivel también proporciona muchos otros beneficios convenientes. En particular, la API se encarga de guardar y restablecer los puntos de control del modelo, de manera que puedas reanudar un trabajo de entrenamiento interrumpido en el mismo lugar en que se detuvo.

Consulta la lista de operaciones de TensorFlow disponibles en Cloud TPU.

Cómo configurar una máquina de TPU personalizada

Un trabajo de entrenamiento de TPU se ejecuta en una configuración de VM doble. Una VM (la instancia principal) ejecuta el código de Python. La instancia principal controla el servidor de TensorFlow que se ejecuta en un trabajador TPU.

A fin de usar una TPU con AI Platform, configura el trabajo de entrenamiento para acceder a una máquina con TPU habilitada de una de las tres maneras que se indican a continuación:

  • Usa el nivel de escala BASIC_TPU. Puedes usar este método para acceder a los aceleradores TPU v2.
  • Usa un trabajador cloud_tpu y un tipo de máquina heredado para la VM principal. Puedes usar este método para acceder a los aceleradores TPU v2.
  • Usa un trabajador cloud_tpu y un tipo de máquina de Compute Engine para la VM principal. Puedes usar este método para acceder a los aceleradores TPU v2 o TPU v3. Los aceleradores TPU v3 están disponibles en versión Beta.

Máquina básica habilitada para TPU

Configura el nivel de escala en BASIC_TPU para obtener una VM principal y una VM de TPU, incluida una TPU con ocho núcleos TPU v2, como hiciste al momento de ejecutar la muestra anterior.

Usa el trabajador TPU en una configuración de tipo de máquina heredada

De manera alternativa, puedes definir una configuración de máquina personalizada si necesitas más recursos de procesamiento en la VM principal:

  • Configura el nivel de escala en CUSTOM.
  • Configura la VM maestra para usar un tipo de máquina heredado que se adapte a tus requisitos de trabajo.
  • Establece workerType en cloud_tpu para obtener una VM de TPU, incluida una Cloud TPU con ocho núcleos TPU v2.
  • Configura workerCount en 1.
  • No especifiques un servidor de parámetros cuando utilices Cloud TPU. Si parameterServerCount es mayor que cero, el servicio rechaza la solicitud de trabajo.

En el siguiente ejemplo, se muestra un archivo config.yaml que usa este tipo de configuración:

trainingInput:
  scaleTier: CUSTOM
  masterType: complex_model_m
  workerType: cloud_tpu
  workerCount: 1

Usa el trabajador TPU en una configuración de tipo de máquina de Compute Engine

También puedes definir una configuración de máquina personalizada con un tipo de máquina de Compute Engine para la VM principal y una acceleratorConfig conectada a la VM de TPU.

Puedes usar este tipo de configuración para configurar un trabajador TPU con ocho núcleos TPU v2 (similar a una configuración sin acceleratorConfig) o un trabajador TPU con ocho núcleos TPU v3 (Beta). Lee más sobre la diferencia entre los aceleradores TPU v2 y TPU v3.

El uso de un tipo de máquina de Compute Engine también proporciona más flexibilidad para configurar tu VM principal:

  • Configura el nivel de escala en CUSTOM.
  • Configura la VM principal para usar un tipo de máquina de Compute Engine que se adapte a los requisitos del trabajo.
  • Configura workerType en cloud_tpu.
  • Agrega una workerConfig con un campo acceleratorConfig. Dentro de esa acceleratorConfig, establece type en TPU_V2, o bien TPU_V3 y count en 8. No puedes adjuntar ningún otro número de núcleos de TPU.
  • Configura workerCount en 1.
  • No especifiques un servidor de parámetros cuando utilices Cloud TPU. Si parameterServerCount es mayor que cero, el servicio rechaza la solicitud de trabajo.

En el siguiente ejemplo, se muestra un archivo config.yaml que usa este tipo de configuración:

TPU v2

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V2
      count: 8

TPU v3 (versión Beta)

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V3
      count: 8

Usa un contenedor personalizado en un trabajador TPU

Si deseas ejecutar un contenedor personalizado en el trabajador TPU en lugar de usar una de las versiones del entorno de ejecución de AI Platform que admiten TPU, debes especificar un campo de configuración adicional cuando envíes el trabajo de entrenamiento. Configura la tpuTfVersion en una versión de entorno de ejecución que incluya la misma versión de TensorFlow que tu utiliza tu contenedor. Debes especificar una versión del entorno de ejecución que se admita actualmente para el entrenamiento con TPU.

Debido a que estás configurando el trabajo para usar un contenedor personalizado, AI Platform Training no utiliza el entorno de esta versión del entorno de ejecución cuando ejecuta el trabajo de entrenamiento. Sin embargo, AI Platform Training solicita este campo a fin de poder preparar de forma adecuada el trabajador TPU para la versión de TensorFlow que utiliza el contenedor personalizado.

En el siguiente ejemplo, se muestra un archivo config.yaml con una configuración de TPU similar a la que se utilizó en la sección anterior, excepto que, en este caso, la VM principal y el trabajador TPU ejecutan diferentes contenedores personalizados:

TPU v2

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/your-project-id/your-master-image-name:your-master-tag-name
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/your-project-id/your-worker-image-name:your-worker-tag-name
    tpuTfVersion: 1.14
    acceleratorConfig:
      type: TPU_V2
      count: 8

TPU v3 (versión Beta)

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/your-project-id/your-master-image-name:your-master-tag-name
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/your-project-id/your-worker-image-name:your-worker-tag-name
    tpuTfVersion: 1.14
    acceleratorConfig:
      type: TPU_V3
      count: 8

Si usas el comando gcloud beta ai-platform jobs submit training para enviar tu trabajo de entrenamiento, puedes especificar el campo tpuTfVersion de la API con la marca --tpu-tf-version en vez de hacerlo en un archivo config.yaml.

Obtén más información sobre el entrenamiento distribuido con contenedores personalizados.

Próximos pasos