Usa la 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 trabajos de entrenamiento de Cloud Machine Learning Engine con Cloud TPU. Cloud ML Engine proporciona una interfaz de administración de trabajo, de manera que no tienes que administrar las TPU por ti mismo. En su lugar, puedes usar la API de jobs de Cloud ML Engine del mismo modo que la usas para entrenar en una CPU o GPU.

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

Configura y prueba el entorno de GCP

Configura el entorno de GCP revisando 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 al nombre de la cuenta de servicio de Cloud TPU asociada con el proyecto de GCP:

  1. Obtén el nombre de la cuenta de servicio de Cloud TPU llamando a projects.getConfig. Por 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 siguientes pasos en Google Cloud Platform Console o con el comando de gcloud:

Console

  1. Accede a Google Cloud Platform Console y elige el proyecto para 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 servicios 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 de ml.serviceAgent a la cuenta de servicios 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

En esta sección, se muestra cómo entrenar el modelo Tensorflow ResNet-50 de referencia con un conjunto de datos falso proporcionado 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 en Cloud ML Engine:

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

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

    cd tpu-r1.12/models/official/
    
  3. Edita ./resnet/resnet_main.py y cambia al código para usar importaciones relativas explícitas al momento de importar 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 proporcionan actualmente acceso a las TPU:

    • us-central1

  6. Envía el trabajo de entrenamiento con el comando gcloud ml-engine jobs submit training:

    gcloud ml-engine jobs submit training $JOB_NAME \
            --staging-bucket $STAGING_BUCKET \
            --runtime-version 1.12 \
            --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 Cloud ML Engine 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 proporcionan actualmente acceso a las TPU:

  • us-central1

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

Control de versiones de TensorFlow y Cloud ML Engine

Las versiones del entorno de ejecución 1.9, 1.11 y 1.12 de Cloud ML Engine se encuentran disponibles para entrenar los modelos en Cloud TPU. Consulta más información sobre las versiones del entorno de ejecución de Cloud ML Engine y las versiones correspondientes de TensorFlow.

En Beta, 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 usa en el código de entrenamiento.

Configura 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.

Para usar una TPU con Cloud ML Engine, configura el trabajo de entrenamiento con el fin de que tenga acceso a una máquina con TPU habilitada.

  • Puedes configurar el nivel de escala en BASIC_TPU para obtener una instancia principal de VM y una VM de TPU que incluya una TPU, como hiciste al momento de ejecutar el ejemplo anterior.
  • De manera alternativa, puedes definir una configuración de máquina personalizada si necesitas más recursos de procesamiento en la instancia principal de VM:

    • Configura el nivel de escala en CUSTOM.
    • Configura el tipo de máquina de la instancia principal para que se ajuste a los requisitos del trabajo.
    • Configura la tarea de trabajo para que use un tipo de máquina cloud_tpu, de este modo obtienes una VM de TPU, incluida una Cloud TPU.
    • Ten en cuanta que no debes especificar un servidor de parámetro cuando usas Cloud TPU. Si parameterServerCount es mayor que cero, el servicio rechaza la solicitud de trabajo.
    • workerCount debe ser 1.

Consulta más información sobre el nivel de escala y el tipo de máquina.

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 de una máquina, debes usar la API de TensorFlow TPUEstimator, que se hereda de la API de TensorFlow Estimator de alto nivel.

  • TPUEstimator maneja muchos de los detalles de la ejecución en dispositivos TPU, como la replicación de entradas y modelos para 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 puntos de control, de manera que puedas reanudar un trabajo de entrenamiento interrumpido en el mismo punto en que se detuvo.

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

¿Qué sigue?

¿Te sirvió esta página? Envíanos tu opinión:

Enviar comentarios sobre…

Cloud ML Engine para TensorFlow