Usa las TPU para entrenar tu modelo

Las unidades de procesamiento tensorial (TPU) son los ASIC de Google personalizados que se usan para acelerar las cargas de trabajo de aprendizaje automático. Puedes ejecutar tus trabajos de entrenamiento en AI Platform Training con Cloud TPU. AI Platform Training proporciona una interfaz de administración de trabajos para que no tengas que administrar la TPU por tu cuenta. En su lugar, puedes usar la API jobs de AI Platform Training de la misma forma que la usas para entrenar 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 tu entorno de Google Cloud

Para configurar tu entorno de Google Cloud, revisa 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 asociada con tu proyecto de Google Cloud:

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

   PROJECT_ID=PROJECT_ID
   
   curl -H "Authorization: Bearer $(gcloud auth print-access-token)"  \
       https://ml.googleapis.com/v1/projects/$PROJECT_ID:getConfig
   

2. Guarda el valor de los campos serviceAccountProject y tpuServiceAccount que muestra la API.

3. Inicializa la cuenta de servicio de Cloud TPU:

   curl -H "Authorization: Bearer $(gcloud auth print-access-token)"  \
     -H "Content-Type: application/json" -d '{}'  \
     https://serviceusage.googleapis.com/v1beta1/projects/<serviceAccountProject>/services/tpu.googleapis.com:generateServiceIdentity
   

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 la consola de Google Cloud o usa el comando gcloud:

Console

1. Accede a la consola de Google Cloud y elige el proyecto en el que usas la TPU.
2. Elige IAM y administración > 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 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=PROJECT_ID
   SVC_ACCOUNT=[e-mail]
   

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
   

Si quieres obtener más detalles para otorgar funciones a las cuentas de servicio, consulta la documentación de IAM.

Ejemplo: Entrena un modelo de MNIST de muestra

En esta sección, se muestra cómo entrenar un modelo MNIST de muestra con una TPU y la versión 2.11 del entorno de ejecución. 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.

En este ejemplo, se da por sentado que usas una shell Bash con la CLI de gcloud instalada. Ejecuta los siguientes comandos para obtener el código y enviar tu trabajo de entrenamiento a AI Platform Training:

1. Descarga el código para los modelos de referencia de TensorFlow y navega al directorio con el código de muestra:

   git clone https://github.com/tensorflow/models.git \
     --branch=v2.11.0 \
     --depth=1
   
   cd models
   

2. Crea un archivo setup.py en el directorio models. Esto garantiza que el comando gcloud ai-platform jobs submit training incluya todos los subpaquetes necesarios dentro del directorio models/official cuando cree un tarball del código de entrenamiento y asegure que AI Platform Training instale TensorFlow Datasets como una dependencia cuando ejecuta el trabajo de entrenamiento. Este código de entrenamiento depende de los conjuntos de datos de TensorFlow para cargar los datos de MNIST.

   Para crear el archivo setup.py, ejecuta el siguiente comando en la shell:

   cat << END > setup.py
   from setuptools import find_packages
   from setuptools import setup
   
   setup(
       name='official',
       install_requires=[
          'tensorflow-datasets~=3.1',
          'tensorflow-model-optimization>=0.4.1'
      ],
       packages=find_packages()
   )
   END
   

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

   gcloud ai-platform jobs submit training tpu_mnist_1 \
     --staging-bucket=gs://BUCKET_NAME \
     --package-path=official \
     --module-name=official.vision.image_classification.mnist_main \
     --runtime-version=2.11 \
     --python-version=3.7 \
     --scale-tier=BASIC_TPU \
     --region=us-central1 \
     -- \
     --distribution_strategy=tpu \
     --data_dir=gs://tfds-data/datasets \
     --model_dir=gs://BUCKET_NAME/tpu_mnist_1_output
   

   Reemplaza BUCKET_NAME por el nombre de un depósito de Cloud Storage en tu proyecto de Google Cloud. La CLI de gcloud sube el código de entrenamiento empaquetado a este bucket, y AI Platform Training guarda el resultado del entrenamiento en el bucket.

4. Supervisa tu trabajo de entrenamiento. Cuando se completa el trabajo, puedes ver su resultado en el directorio gs://BUCKET_NAME/tpu_mnist_1_output.

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

En esta sección, se brinda más información sobre cómo configurar un trabajo y cómo entrenar un modelo en AI Platform Training 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
• europe-west4

A fin de conocer a fondo las regiones disponibles para los servicios de AI Platform Training, 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 Training

Versiones del entorno de ejecución de AI Platform Training 1.15, 2.1, 2.2, 2.3, 2.4, 2.5, 2.6, 2.7, 2.8, 2.9 y 2.11 están disponibles para entrenar tus modelos Cloud TPU Obtén más información sobre las versiones del entorno de ejecución de AI Platform Training y las versiones de TensorFlow correspondientes.

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 al servidor gRPC de la TPU

En el programa de TensorFlow, usa TPUClusterResolver para conectarte con el servidor gRPC de la TPU que se ejecuta en la VM de TPU.

En la guía de TensorFlow para usar TPU, se muestra cómo usar TPUClusterResolver con la estrategia de distribución TPUStrategy.

Sin embargo, debes realizar un cambio importante cuando usas TPUClusterResolver para el código que se ejecuta en AI Platform Training: no proporciones ningún argumento cuando crees la instancia TPUClusterResolver. Cuando los argumentos de palabra clave tpu, zone y project se configuran con su valor predeterminado de None, AI Platform Training proporciona de forma automática al agente de resolución del clúster los detalles de conexión necesarios a través variables de entorno.

En el siguiente ejemplo de TensorFlow 2, se muestra cómo inicializar un agente de resolución del clústeres y una estrategia de distribución para el entrenamiento en AI Platform Training:

import tensorflow as tf

resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

Usa las TPU en el código de TensorFlow

Para usar las TPU en una máquina, usa la API de TPUStrategy de TensorFlow. En la guía de TensorFlow para usar TPU, se muestra cómo hacerlo.

Para entrenar con las TPU en TensorFlow 1, puedes usar la API de TPUEstimator en su lugar. En la guía de Cloud TPU para la API de TPUEstimator, se muestra cómo hacerlo.

En la documentación de Cloud TPU, también se proporciona una lista de operaciones de TensorFlow de bajo nivel disponibles en Cloud TPU.

Usa las TPU en el código PyTorch

Para usar una TPU cuando usas un contenedor de PyTorch compilado con anterioridad, usa el paquete torch_xla. Obtén más información sobre cómo usar torch_xla para TPU en el entrenamiento de la documentación de PyTorch. Para obtener más ejemplos sobre el uso de torch_xla, consulta los instructivos en el repositorio de GitHub de PyTorch XLA.

Ten en cuenta que cuando entrenas con una TPU en AI Platform Training, usas un solo dispositivo XLA, no con varios dispositivos de XLA.

Consulta también la siguiente sección en esta página sobre cómo configurar tu trabajo de entrenamiento para PyTorch y TPU.

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 principal controla el servidor de TensorFlow que se ejecuta en un trabajador TPU.

A fin de usar una TPU con AI Platform Training, configura el trabajo de entrenamiento para que acceda a una máquina compatible con TPU de tres maneras:

• 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 de 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.

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:

• Establece 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 uses 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

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 la VM principal:

• Establece 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.
• Establece workerType en cloud_tpu.
• Agrega una workerConfig con un campo acceleratorConfig. Dentro de acceleratorConfig, establece type en TPU_V2 o TPU_V3, y count en 8. No puedes adjuntar ningún otro número de núcleos de TPU.
• Establece workerCount en 1.
• No especifiques un servidor de parámetros cuando uses 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 (Beta)

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

Usa pods de TPU

Vista previa

Esta función está sujeta a las "Condiciones de las Ofertas de Fase Previa a la DG" de la sección Condiciones Generales del Servicio de las Condiciones Específicas del Servicio. Las funciones de fase previa a la DG están disponibles “tal cual” y podrían tener asistencia limitada. Para obtener más información, consulta las descripciones de la etapa de lanzamiento.

Un pod de TPU es una colección de dispositivos de TPU conectados por interfaces de red dedicadas de alta velocidad. Un pod de TPU puede tener hasta 2,048 núcleos de TPU, lo que te permite distribuir la carga de procesamiento en varias TPU.

Para usar pods de TPU, primero debes enviar una solicitud de aumento de cuota.

En los siguientes archivos config.yaml de ejemplo, se muestra cómo usar pods de TPU:

Pods de TPU v2

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

Pods de TPU v3

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

Existen limitaciones en la cantidad de núcleos de pod que se pueden usar para cada tipo de TPU. Opciones de configuración disponibles:

Tipo de pod de TPU	Cantidad de núcleos de pod disponibles para usar
TPU_V2_POD	32, 128, 256, 512
TPU_V3_POD	32, 128, 256

Para obtener más detalles sobre cómo usar todos los núcleos de los pods de TPU, consulta la documentación de Cloud TPU sobre los pods de TPU.

Usa un contenedor PyTorch compilado con anterioridad en un trabajador TPU

Si deseas realizar un entrenamiento de PyTorch con una TPU, debes especificar el campo tpuTfVersion en el trainingInput de tu trabajo de entrenamiento. Configura tpuTfVersion para que coincida con la versión del contenedor de PyTorch compilado previamente que se usa para el entrenamiento.

AI Platform Training admite el entrenamiento con TPU para los siguientes contenedores de PyTorch preconfigurados:

URI de la imagen de contenedor	tpuTfVersion
gcr.io/cloud-ml-public/training/pytorch-xla.1-11	pytorch-1.11
gcr.io/cloud-ml-public/training/pytorch-xla.1-10	pytorch-1.10
gcr.io/cloud-ml-public/training/pytorch-xla.1-9	pytorch-1.9
gcr.io/cloud-ml-public/training/pytorch-xla.1-7	pytorch-1.7
gcr.io/cloud-ml-public/training/pytorch-xla.1-6	pytorch-1.6

Por ejemplo, puedes usar el archivo config.yaml para entrenar con el contenedor compilado previamente de PyTorch 1.11 a fin de configurar el entrenamiento:

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-11
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-11
    tpuTfVersion: pytorch-1.11
    acceleratorConfig:
      type: TPU_V2
      count: 8

Consulta también la sección anterior de esta página Usa de TPU en el código de PyTorch.

Usa un contenedor personalizado en un trabajador TPU

Beta

Esta función está sujeta a las "Condiciones de las Ofertas de Fase Previa a la DG" de la sección Condiciones Generales del Servicio de las Condiciones Específicas del Servicio. Las funciones de fase previa a la DG están disponibles “tal cual” y podrían tener asistencia limitada. Para obtener más información, consulta las descripciones de la etapa de lanzamiento.

Si deseas ejecutar un contenedor personalizado en tu trabajador TPU en lugar de usar una de las versiones del entorno de ejecución de AI Platform Training compatibles con TPU, debes especificar un campo de configuración adicional cuando envíes el trabajo de entrenamiento. Configura la tpuTfVersion en una versión del entorno de ejecución que incluya la misma versión de TensorFlow que tu usa 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 usó 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: 2.11
    acceleratorConfig:
      type: TPU_V2
      count: 8

TPU v3 (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: 2.11
    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.

Usa TPUClusterResolver después de aprovisionar la TPU

Cuando usas un contenedor personalizado, debes esperar a que se aprovisione la TPU antes de poder llamar a TPUClusterResolver para usarla. En el siguiente código de muestra, podrás ver cómo manejar la lógica TPUClusterResolver:

def wait_for_tpu_cluster_resolver_ready():
  """Waits for `TPUClusterResolver` to be ready and return it.

  Returns:
    A TPUClusterResolver if there is TPU machine (in TPU_CONFIG). Otherwise,
    return None.
  Raises:
    RuntimeError: if failed to schedule TPU.
  """
  tpu_config_env = os.environ.get('TPU_CONFIG')
  if not tpu_config_env:
    tf.logging.info('Missing TPU_CONFIG, use CPU/GPU for training.')
    return None

  tpu_node = json.loads(tpu_config_env)
  tf.logging.info('Waiting for TPU to be ready: \n%s.', tpu_node)

  num_retries = 40
  for i in range(num_retries):
    try:
      tpu_cluster_resolver = (
          tf.contrib.cluster_resolver.TPUClusterResolver(
              tpu=[tpu_node['tpu_node_name']],
              zone=tpu_node['zone'],
              project=tpu_node['project'],
              job_name='worker'))
      tpu_cluster_resolver_dict = tpu_cluster_resolver.cluster_spec().as_dict()
      if 'worker' in tpu_cluster_resolver_dict:
        tf.logging.info('Found TPU worker: %s', tpu_cluster_resolver_dict)
        return tpu_cluster_resolver
    except Exception as e:
      if i < num_retries - 1:
        tf.logging.info('Still waiting for provisioning of TPU VM instance.')
      else:
        # Preserves the traceback.
        raise RuntimeError('Failed to schedule TPU: {}'.format(e))
    time.sleep(10)

  # Raise error when failed to get TPUClusterResolver after retry.
  raise RuntimeError('Failed to schedule TPU.')

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

¿Qué sigue?

• Obtén más información sobre los modelos de entrenamiento en AI Platform Training.
• Obtén información sobre el ajuste de hiperparámetros en AI Platform Training y pon especial atención a los detalles del ajuste de hiperparámetros con Cloud TPU.
• Explora modelos de referencia adicionales para Cloud TPU.
• Optimiza los modelos de Cloud TPU con las recomendaciones de Cloud TPU.
• Consulta la solución de problemas y Preguntas frecuentes de Cloud TPU para obtener ayuda con el diagnóstico y la resolución de problemas.