Vertex AI: Ajuste de hiperparámetros distribuido

1. Resumen

En este lab, aprenderás a usar Vertex AI para el ajuste de hiperparámetros y el entrenamiento distribuido. Si bien en este lab se usa TensorFlow para el código del modelo, los conceptos también se aplican a otros frameworks de AA.

Qué aprenderás

Aprenderás a hacer lo siguiente:

• Entrenar un modelo con entrenamiento distribuido en un contenedor personalizado
• Iniciar varias pruebas de tu código de entrenamiento para el ajuste automático de hiperparámetros

El costo total de la ejecución de este lab en Google Cloud es de aproximadamente $6.

2. Introducción a Vertex AI

En este lab, se utiliza la oferta de productos de IA más reciente de Google Cloud. Vertex AI integra las ofertas de AA de Google Cloud en una experiencia de desarrollo fluida. Anteriormente, se podía acceder a los modelos personalizados y a los entrenados con AutoML mediante servicios independientes. La nueva oferta combina ambos en una sola API, junto con otros productos nuevos. También puede migrar proyectos existentes a Vertex AI. Para enviarnos comentarios, visite la página de asistencia.

Vertex AI incluye muchos productos distintos para respaldar flujos de trabajo de AA de extremo a extremo. Este lab se centrará en el entrenamiento y en Workbench.

3. Descripción general del caso de uso

En este lab, utilizarás el ajuste de hiperparámetros a fin de descubrir los parámetros óptimos para un modelo de clasificación de imágenes entrenado con el conjunto de datos de caballos o seres humanos de TensorFlow.

Ajuste de hiperparámetros

El ajuste de hiperparámetros con Vertex AI Training funciona mediante la ejecución de diversas pruebas de tu aplicación de entrenamiento con valores para tus hiperparámetros elegidos, que se establecen dentro de los límites que especifiques. Vertex AI realiza un seguimiento de los resultados de cada prueba y aplica ajustes en las pruebas posteriores.

Para usar el ajuste de hiperparámetros con el entrenamiento de Vertex AI, debes realizar los siguientes dos cambios en tu código de entrenamiento:

1. Define un argumento de la línea de comandos en el módulo de entrenamiento principal para cada hiperparámetro que desees ajustar.
2. Usa el valor que pasaste en esos argumentos para configurar el hiperparámetro correspondiente en el código de tu aplicación.

Entrenamiento distribuido

Si tienes una sola GPU, TensorFlow usará este acelerador para incrementar la velocidad del entrenamiento de modelos sin que debas realizar ninguna acción adicional. Sin embargo, si quieres ir más rápido mediante el uso de varias GPU, deberás usar tf.distribute, que es el módulo de TensorFlow para procesamiento en varios dispositivos.

En este lab, se usa tf.distribute.MirroredStrategy, que puedes agregar a tus aplicaciones de entrenamiento solo con unos pocos cambios en el código. Esta estrategia crea una copia del modelo en cada GPU de tu máquina. Las actualizaciones de gradientes se realizarán de forma síncrona. Esto significa que cada GPU procesa los modelos de propagación y retropropagación en el modelo en una porción diferente de los datos de entrada. Luego, los gradientes procesados de cada una de estas porciones se agregan en todas las GPU y se promedian en un proceso conocido como all‑reduce. Los parámetros del modelo se actualizan con estos gradientes promediados.

No es necesario que sepas los detalles para completar este lab, pero si deseas obtener más información sobre cómo funciona el entrenamiento distribuido en TensorFlow, mira el siguiente video:

4. Configura el entorno

Para ejecutar este codelab, necesitarás un proyecto de Google Cloud Platform que tenga habilitada la facturación. Para crear un proyecto, sigue estas instrucciones.

Paso 1: Habilita la API de Compute Engine

Ve a Compute Engine y selecciona Habilitar (si aún no está habilitada).

Paso 2: Habilita la API de Container Registry

Navega a Container Registry y selecciona Habilitar si aún no lo has hecho. Deberás usarla con el fin de crear un contenedor para tu trabajo de entrenamiento personalizado.

Paso 3: Habilita la API de Vertex AI

Navega hasta la sección Vertex AI en Cloud Console y haz clic en Habilitar API de Vertex AI (Enable Vertex AI API).

Paso 4: Crea una instancia de Vertex AI Workbench

En la sección Vertex AI de Cloud Console, haz clic en Workbench:

Habilita la API de Notebooks si aún no está habilitada.

Una vez habilitada, haz clic en NOTEBOOKS ADMINISTRADOS (MANAGED NOTEBOOKS):

Luego, selecciona NUEVO NOTEBOOK (NEW NOTEBOOK).

Asígnale un nombre al notebook y, luego, haz clic en Configuración avanzada (Advanced settings).

En Configuración avanzada (Advanced settings), activa la opción Habilitar el cierre inactivo (Enable Idle Shutdown) y establece la cantidad de minutos en 60. Esto provocará que el notebook se cierre automáticamente cuando no esté en uso para que no se generen costos innecesarios.

En Seguridad (Security), selecciona la opción “Habilitar terminal” (Enable terminal) si aún no está habilitada.

Puedes dejar el resto de la configuración avanzada tal como está.

Luego, haz clic en Crear. La instancia tardará algunos minutos en aprovisionarse.

Una vez que se crea la instancia, selecciona Abrir JupyterLab (Open JupyterLab).

La primera vez que uses una instancia nueva, se te solicitará que te autentiques. Sigue los pasos en la IU para hacerlo.

5. Escribe el código de entrenamiento

Para comenzar, en el menú Launcher, abre una ventana de terminal en tu instancia de notebook:

Crea un directorio nuevo llamado vertex-codelab y ábrelo con el comando cd.

mkdir vertex-codelab
cd vertex-codelab

Ejecuta el siguiente comando a fin de crear un directorio para el código de entrenamiento y un archivo de Python en el que agregarás el código:

mkdir trainer
touch trainer/task.py

Ahora, deberías tener lo siguiente en el directorio vertex-codelab:

+ trainer/
    + task.py

A continuación, abre el archivo task.py que acabas de crear y pega todo el siguiente código.

import tensorflow as tf
import tensorflow_datasets as tfds
import argparse
import hypertune
import os

NUM_EPOCHS = 10
BATCH_SIZE = 64

def get_args():
  '''Parses args. Must include all hyperparameters you want to tune.'''

  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--learning_rate',
      required=True,
      type=float,
      help='learning rate')
  parser.add_argument(
      '--momentum',
      required=True,
      type=float,
      help='SGD momentum value')
  parser.add_argument(
      '--num_units',
      required=True,
      type=int,
      help='number of units in last hidden layer')
  args = parser.parse_args()
  return args

def preprocess_data(image, label):
  '''Resizes and scales images.'''

  image = tf.image.resize(image, (150,150))
  return tf.cast(image, tf.float32) / 255., label

def create_dataset(batch_size):
  '''Loads Horses Or Humans dataset and preprocesses data.'''

  data, info = tfds.load(name='horses_or_humans', as_supervised=True, with_info=True)

  # Create train dataset
  train_data = data['train'].map(preprocess_data)
  train_data  = train_data.shuffle(1000)
  train_data  = train_data.batch(batch_size)

  # Create validation dataset
  validation_data = data['test'].map(preprocess_data)
  validation_data  = validation_data.batch(batch_size)

  return train_data, validation_data

def create_model(num_units, learning_rate, momentum):
  '''Defines and compiles model.'''

  inputs = tf.keras.Input(shape=(150, 150, 3))
  x = tf.keras.layers.Conv2D(16, (3, 3), activation='relu')(inputs)
  x = tf.keras.layers.MaxPooling2D((2, 2))(x)
  x = tf.keras.layers.Conv2D(32, (3, 3), activation='relu')(x)
  x = tf.keras.layers.MaxPooling2D((2, 2))(x)
  x = tf.keras.layers.Conv2D(64, (3, 3), activation='relu')(x)
  x = tf.keras.layers.MaxPooling2D((2, 2))(x)
  x = tf.keras.layers.Flatten()(x)
  x = tf.keras.layers.Dense(num_units, activation='relu')(x)
  outputs = tf.keras.layers.Dense(1, activation='sigmoid')(x)
  model = tf.keras.Model(inputs, outputs)
  model.compile(
      loss='binary_crossentropy',
      optimizer=tf.keras.optimizers.SGD(learning_rate=learning_rate, momentum=momentum),
      metrics=['accuracy'])
  return model

def main():
  args = get_args()

  # Create distribution strategy
  strategy = tf.distribute.MirroredStrategy()

  # Get data
  GLOBAL_BATCH_SIZE = BATCH_SIZE * strategy.num_replicas_in_sync
  train_data, validation_data = create_dataset(GLOBAL_BATCH_SIZE)

  # Wrap variable creation within strategy scope
  with strategy.scope():
    model = create_model(args.num_units, args.learning_rate, args.momentum)

  # Train model
  history = model.fit(train_data, epochs=NUM_EPOCHS, validation_data=validation_data)

  # Define metric
  hp_metric = history.history['val_accuracy'][-1]

  hpt = hypertune.HyperTune()
  hpt.report_hyperparameter_tuning_metric(
      hyperparameter_metric_tag='accuracy',
      metric_value=hp_metric,
      global_step=NUM_EPOCHS)

if __name__ == "__main__":
    main()

Analicemos con más detalle el código y examinemos los componentes específicos del entrenamiento distribuido y el ajuste de hiperparámetros.

Entrenamiento distribuido

1. En la función main(), se crea el objeto MirroredStrategy. A continuación, agrega la creación de las variables de tu modelo al alcance de la estrategia. En este paso, se indica a TensorFlow qué variables se deben duplicar en las GPU.
2. El tamaño del lote se escala verticalmente según num_replicas_in_sync. Escalar el tamaño del lote es una práctica recomendada cuando se usan estrategias de paralelismo de datos síncronos en TensorFlow. Puedes obtener más información aquí.

Ajuste de hiperparámetros

1. La secuencia de comandos importa la biblioteca hypertune. Más adelante, cuando compilemos la imagen del contenedor, tendremos que asegurarnos de instalar esta biblioteca.
2. La función get_args() define un argumento de la línea de comandos para cada hiperparámetro que deseas ajustar. En este ejemplo, los hiperparámetros que se ajustarán son la tasa de aprendizaje, el valor de impulso en el optimizador y la cantidad de unidades en la última capa oculta del modelo, pero puedes experimentar con otros. El valor que se pasó en esos argumentos se usa para establecer el hiperparámetro correspondiente en el código (p. ej., establecer learning_rate = args.learning_rate)
3. Al final de la función main(), se usa la biblioteca hypertune para definir la métrica que quieres optimizar. En TensorFlow, el método model.fit de Keras muestra un objeto History. El atributo History.history es un registro de los valores de pérdida de entrenamiento y de valores de métricas en ciclos sucesivos. Si pasas datos de validación a model.fit, el atributo History.history también incluirá valores de métricas y pérdida de validación. Por ejemplo, si entrenaste un modelo durante tres ciclos con datos de validación y proporcionaste accuracy como métrica, el atributo History.history se verá como el siguiente diccionario.

{
 "accuracy": [
   0.7795261740684509,
   0.9471358060836792,
   0.9870933294296265
 ],
 "loss": [
   0.6340447664260864,
   0.16712145507335663,
   0.04546636343002319
 ],
 "val_accuracy": [
   0.3795261740684509,
   0.4471358060836792,
   0.4870933294296265
 ],
 "val_loss": [
   2.044623374938965,
   4.100203514099121,
   3.0728273391723633
 ]

Si quieres que el servicio de ajuste de hiperparámetros descubra qué valores maximizan la exactitud de la validación del modelo, debe definir la métrica como la última entrada (o NUM_EPOCS - 1) de la lista val_accuracy. Luego, pasa esta métrica a una instancia de HyperTune. Puedes elegir la string que desees para hyperparameter_metric_tag, pero deberás volver a usarla más adelante cuando inicies el trabajo de ajuste de hiperparámetros.

6. Aloja el código en contenedores

El primer paso para alojar tu código en contenedores es crear un Dockerfile. En él, incluirás todos los comandos necesarios para ejecutar la imagen. Se instalarán todas las bibliotecas necesarias y se configurará el punto de entrada para el código de entrenamiento.

Paso 1: Escribe el Dockerfile

Desde tu terminal, asegúrate de estar en el directorio vertex-codelab y crea un Dockerfile vacío:

touch Dockerfile

Ahora, deberías tener lo siguiente en el directorio vertex-codelab:

+ Dockerfile
+ trainer/
    + task.py

Abre el Dockerfile y copia lo siguiente:

FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-7

WORKDIR /

# Installs hypertune library
RUN pip install cloudml-hypertune

# Copies the trainer code to the docker image.
COPY trainer /trainer

# Sets up the entry point to invoke the trainer.
ENTRYPOINT ["python", "-m", "trainer.task"]

Este Dockerfile usa la imagen de Docker TensorFlow Enterprise 2.7 GPU de los Contenedores de aprendizaje profundo. Los Contenedores de aprendizaje profundo de Google Cloud tienen instalados varios frameworks típicos del AA y ciencia de datos. Después de descargar la imagen, el Dockerfile configurará el punto de entrada para el código de entrenamiento.

Paso 2: Compila el contenedor

Desde tu terminal, ejecuta lo siguiente a fin de definir una variable de entorno para tu proyecto y asegúrate de reemplazar your-cloud-project con el ID de tu proyecto:

Para obtener el ID del proyecto, ejecuta gcloud config list --format 'value(core.project)' en tu terminal

PROJECT_ID='your-cloud-project'

Define una variable con el URI de la imagen de contenedor en Google Container Registry:

IMAGE_URI="gcr.io/$PROJECT_ID/horse-human-codelab:latest"

Configura el Docker

gcloud auth configure-docker

Luego, ejecuta el siguiente comando para compilar el contenedor desde la raíz de tu directorio vertex-codelab:

docker build ./ -t $IMAGE_URI

Por último, envíala a Google Container Registry:

docker push $IMAGE_URI

Paso 3: Crea un bucket de Cloud Storage

En nuestro trabajo de entrenamiento, pasaremos la ruta de acceso a un bucket de etapa de pruebas.

Ejecuta el siguiente comando en tu terminal para crear un bucket nuevo en tu proyecto.

BUCKET_NAME="gs://${PROJECT_ID}-hptune-bucket"
gsutil mb -l us-central1 $BUCKET_NAME

7. Inicia el trabajo de ajuste de hiperparámetros

Paso 1: Crea un trabajo de entrenamiento personalizado con ajuste de hiperparámetros

Desde el selector, abre un nuevo notebook de TensorFlow 2.

Importa el SDK de Vertex AI para Python.

from google.cloud import aiplatform
from google.cloud.aiplatform import hyperparameter_tuning as hpt

Para iniciar el trabajo de ajuste de hiperparámetros, primero debes definir el worker_pool_specs, que especifica el tipo de máquina y la imagen de Docker. En las siguientes especificaciones, se define una máquina con dos GPU NVIDIA Tesla V100.

Deberás reemplazar {PROJECT_ID} en el image_uri con tu proyecto.

# The spec of the worker pools including machine type and Docker image
# Be sure to replace PROJECT_ID in the "image_uri" with your project.

worker_pool_specs = [{
    "machine_spec": {
        "machine_type": "n1-standard-4",
        "accelerator_type": "NVIDIA_TESLA_V100",
        "accelerator_count": 2
    },
    "replica_count": 1,
    "container_spec": {
        "image_uri": "gcr.io/{PROJECT_ID}/horse-human-codelab:latest"
    }
}]

A continuación, define parameter_spec, que es un diccionario que especifica los parámetros que deseas optimizar. La clave del diccionario es la string que asignaste al argumento de la línea de comandos para cada hiperparámetro, y el valor del diccionario es la especificación del parámetro.

En cada hiperparámetro, debes definir el tipo y los límites de los valores que intentará el servicio de ajuste. Los hiperparámetros pueden ser de tipo doble, de número entero, categóricos o discretos. Si seleccionas el tipo doble o número entero, deberás proporcionar un valor mínimo y uno máximo. Además, si seleccionas categórico o discreto, debes proporcionar los valores. Para los tipos doble y número entero, también deberás proporcionar el valor de escalamiento. Puedes obtener más información acerca de cómo elegir la mejor escala en este video.

# Dictionary representing parameters to optimize.
# The dictionary key is the parameter_id, which is passed into your training
# job as a command line argument,
# And the dictionary value is the parameter specification of the metric.
parameter_spec = {
    "learning_rate": hpt.DoubleParameterSpec(min=0.001, max=1, scale="log"),
    "momentum": hpt.DoubleParameterSpec(min=0, max=1, scale="linear"),
    "num_units": hpt.DiscreteParameterSpec(values=[64, 128, 512], scale=None)
}

La especificación final que se debe definir es metric_spec, que es un diccionario que representa la métrica que se debe optimizar. La clave del diccionario es el hyperparameter_metric_tag que estableces en el código de tu aplicación de entrenamiento y el valor es el objetivo de optimización.

# Dicionary representing metrics to optimize.
# The dictionary key is the metric_id, which is reported by your training job,
# And the dictionary value is the optimization goal of the metric.
metric_spec={'accuracy':'maximize'}

Una vez que se definan las especificaciones, crearás un CustomJob, que es la especificación común que se usará para ejecutar el trabajo en cada una de las pruebas de ajuste de hiperparámetros.

Deberás reemplazar {YOUR_BUCKET} por el bucket que creaste anteriormente.

# Replace YOUR_BUCKET
my_custom_job = aiplatform.CustomJob(display_name='horses-humans',
                              worker_pool_specs=worker_pool_specs,
                              staging_bucket='gs://{YOUR_BUCKET}')

Luego, crea y ejecuta HyperparameterTuningJob.

hp_job = aiplatform.HyperparameterTuningJob(
    display_name='horses-humans',
    custom_job=my_custom_job,
    metric_spec=metric_spec,
    parameter_spec=parameter_spec,
    max_trial_count=6,
    parallel_trial_count=2,
    search_algorithm=None)

hp_job.run()

Debes tener en cuenta los siguientes argumentos:

• max_trial_count: Deberás establecer un límite superior para la cantidad de pruebas que se ejecutarán en el servicio. Si se realizan más pruebas, se obtendrán mejores resultados, pero habrá un punto en el que los beneficios disminuirán, y las pruebas adicionales a partir de ese punto tendrán un efecto nulo o muy limitado en la métrica que intentas optimizar. Como práctica recomendada, comienza con una cantidad más pequeña de pruebas para tener una idea del impacto que generan los hiperparámetros elegidos antes de escalar verticalmente.
• paralelo_trial_count: Si usas pruebas paralelas, el servicio aprovisiona varios clústeres de procesamiento de entrenamiento. Si las aumentas, se reducirá el tiempo que el trabajo de ajuste de hiperparámetros tarda en ejecutarse; sin embargo, esto puede reducir su eficacia general. Esto se debe a que la estrategia de ajuste predeterminada utiliza los resultados de pruebas anteriores para asignar los valores en las pruebas posteriores.
• search_algorithm: Puedes configurar el algoritmo de búsqueda como cuadrícula, aleatorio o predeterminado (Ninguno). La opción predeterminada aplica la optimización bayesiana para buscar el espacio de valores de hiperparámetros posibles y es el algoritmo recomendado. Obtén más información sobre este algoritmo aquí.

Una vez que se inicie el trabajo, podrás realizar un seguimiento del estado en la IU, en la pestaña TRABAJOS DE AJUSTE DE HIPERPARÁMETROS (HYPERPARAMETER TUNING JOBS).

Cuando el trabajo se complete, podrás ver y ordenar los resultados de tus pruebas para descubrir la mejor combinación de valores de hiperparámetros.

🎉 ¡Felicitaciones! 🎉

Aprendiste a usar Vertex AI para hacer lo siguiente:

• Ejecutar un trabajo de ajuste de hiperparámetros con entrenamiento distribuido

Si quieres obtener más información sobre los diferentes componentes de Vertex AI, consulta la documentación.

8. Realiza una limpieza

Debido a que configuramos el notebook para que se agote el tiempo de espera después de 60 minutos de inactividad, no tenemos que preocuparnos por cerrar la instancia. Si quieres cerrar la instancia de forma manual, haz clic en el botón Detener (Stop) en la sección Vertex AI Workbench de la consola. Si quieres borrar el notebook por completo, haz clic en el botón Borrar (Delete).

Para borrar el bucket de almacenamiento, en el menú de navegación de Cloud Console, navega a Almacenamiento, selecciona tu bucket y haz clic en Borrar (Delete):