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:
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
Guarda el valor de los campos
serviceAccountProject
ytpuServiceAccount
que muestra la API.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
- Accede a la consola de Google Cloud y elige el proyecto en el que usas la TPU.
- Elige IAM y administración > IAM.
- Haz clic en el botón Agregar para agregar un miembro al proyecto.
- Ingresa la cuenta de servicio de la TPU en el cuadro de texto Miembros.
- Haz clic en la lista desplegable Funciones.
- Habilita la función Agente de servicios de Cloud ML (Administración de servicios > Agente de servicios de Cloud ML).
gcloud
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=your-tpu-sa-123@your-tpu-sa.google.com.iam.gserviceaccount.com
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:
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
Crea un archivo
setup.py
en el directoriomodels
. Esto garantiza que el comandogcloud ai-platform jobs submit training
incluya todos los subpaquetes necesarios dentro del directoriomodels/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
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.
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
encloud_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
encloud_tpu
. - Agrega una
workerConfig
con un campoacceleratorConfig
. Dentro deacceleratorConfig
, establecetype
enTPU_V2
oTPU_V3
, ycount
en8
. 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
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
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.