Vertex AI: Entrenamiento de varios trabajadores y aprendizaje por transferencia con TensorFlow

1. Descripción general

En este lab, usarás Vertex AI para ejecutar un trabajo de entrenamiento de varios trabajadores para un modelo de TensorFlow.

Qué aprenderá

Aprenderás a hacer lo siguiente:

• Modifica el código de la aplicación de entrenamiento para el entrenamiento de varios trabajadores
• Configura y, luego, inicia un trabajo de entrenamiento de varios trabajadores desde la IU de Vertex AI
• Configura y, luego, inicia un trabajo de entrenamiento de varios trabajadores con el SDK de Vertex

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

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 puedes migrar proyectos existentes a Vertex AI. Para enviarnos comentarios, visita 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 enfocará en los productos que se destacan a continuación: Training y Workbench

3. Descripción general del caso de uso

En este lab, usarás el aprendizaje transferible para entrenar un modelo de clasificación de imágenes en el conjunto de datos de mandioca de TensorFlow Datasets. La arquitectura que usarás es un modelo ResNet50 de la biblioteca tf.keras.applications previamente entrenado en el conjunto de datos de ImageNet.

¿Por qué usar el 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 obtener un aumento adicional con el uso de varias GPUs en una sola máquina o varias (cada una con potencialmente varias GPUs), deberás usar tf.distribute, que es la biblioteca de TensorFlow para ejecutar un procesamiento en varios dispositivos. Un dispositivo hace referencia a una CPU o un acelerador, como GPUs o TPU, en una máquina en la que TensorFlow puede ejecutar operaciones.

La forma más sencilla de comenzar a usar el entrenamiento distribuido es una sola máquina con varios dispositivos GPU. Una estrategia de distribución de TensorFlow del módulo tf.distribute administrará la coordinación de la distribución de datos y las actualizaciones de gradientes en todas las GPUs. Si ya dominas el entrenamiento en un host único y quieres aumentar aún más la escala, agregar varias máquinas a tu clúster puede ayudarte a obtener un aumento del rendimiento aún mayor. Puedes utilizar un clúster de máquinas que sean solo CPU o que tengan una o más GPUs cada una. En este lab, se aborda el último caso y se muestra cómo usar MultiWorkerMirroredStrategy para distribuir el entrenamiento de un modelo de TensorFlow en varias máquinas en Vertex AI.

MultiWorkerMirroredStrategy es una estrategia de paralelismo de datos síncronos que puedes usar con solo algunos cambios en el código. Se crea una copia del modelo en cada dispositivo del clúster. Las actualizaciones de gradientes se realizarán de forma síncrona. Esto significa que cada dispositivo trabajador calcula la propagación y la retropropagación en el modelo en una porción diferente de los datos de entrada. Luego, los gradientes calculados de cada una de estas porciones se agregan en todos los dispositivos de una máquina y en todas las máquinas del clúster, y se reducen (por lo general, un promedio) en un proceso conocido como reducción total. Luego, el optimizador realiza las actualizaciones de parámetros con estos gradientes reducidos, lo que mantiene sincronizados los dispositivos. Para obtener más información sobre el entrenamiento distribuido con 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). La necesitarás para crear la instancia de notebook.

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. Aloja en un contenedor el código de entrenamiento de la aplicación

Para enviar este trabajo de entrenamiento a Vertex, deberás colocar el código de entrenamiento de la aplicación en un contenedor de Docker y enviarlo a Google Container Registry. Mediante este enfoque, puedes entrenar un modelo compilado con cualquier framework.

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

Crea un directorio nuevo llamado cassava y ábrelo con el comando cd:

mkdir cassava
cd cassava

Paso 1: Crea un Dockerfile

El primer paso para alojar el código en un contenedor 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.

En la terminal, crea un Dockerfile vacío:

touch Dockerfile

Abre el Dockerfile y copia lo siguiente:

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

WORKDIR /

# 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 esa imagen, este Dockerfile configura el punto de entrada para el código de entrenamiento. Aún no creas estos archivos; en el siguiente paso, agregarás el código para entrenar y ajustar el modelo.

Paso 2: Crea un bucket de Cloud Storage

En este trabajo de entrenamiento, exportarás el modelo de TensorFlow entrenado a un bucket de Cloud Storage. Desde tu terminal, ejecuta lo siguiente para 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'

Luego, ejecuta el siguiente comando en tu terminal para crear un bucket nuevo en el proyecto.

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

Paso 3: Agrega el código de entrenamiento del modelo

En la terminal, ejecuta el siguiente comando para 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 cassava/:

+ Dockerfile
+ trainer/
    + task.py

A continuación, abre el archivo task.py que acabas de crear y copia el código que aparece más abajo. Deberás reemplazar {your-gcs-bucket} por el nombre del bucket de Cloud Storage que acabas de crear.

import tensorflow as tf
import tensorflow_datasets as tfds
import os


PER_REPLICA_BATCH_SIZE = 64
EPOCHS = 2

# TODO: replace {your-gcs-bucket} with the name of the Storage bucket you created earlier
BUCKET = 'gs://{your-gcs-bucket}/mwms'

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

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


def create_dataset(batch_size):
  '''Loads Cassava dataset and preprocesses data.'''

  data, info = tfds.load(name='cassava', as_supervised=True, with_info=True)
  number_of_classes = info.features['label'].num_classes
  train_data = data['train'].map(preprocess_data,
                                 num_parallel_calls=tf.data.experimental.AUTOTUNE)
  train_data  = train_data.shuffle(1000)
  train_data  = train_data.batch(batch_size)
  train_data  = train_data.prefetch(tf.data.experimental.AUTOTUNE)

  # Set AutoShardPolicy
  options = tf.data.Options()
  options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATA
  train_data = train_data.with_options(options)

  return train_data, number_of_classes


def create_model(number_of_classes):
  '''Creates and compiles pretrained ResNet50 model.'''

  base_model = tf.keras.applications.ResNet50(weights='imagenet', include_top=False)
  x = base_model.output
  x = tf.keras.layers.GlobalAveragePooling2D()(x)
  x = tf.keras.layers.Dense(1016, activation='relu')(x)
  predictions = tf.keras.layers.Dense(number_of_classes, activation='softmax')(x)
  model = tf.keras.Model(inputs=base_model.input, outputs=predictions)

  model.compile(
      loss='sparse_categorical_crossentropy',
      optimizer=tf.keras.optimizers.Adam(0.0001),
      metrics=['accuracy'])

  return model


def _is_chief(task_type, task_id):
  '''Helper function. Determines if machine is chief.'''

  return task_type == 'chief'


def _get_temp_dir(dirpath, task_id):
  '''Helper function. Gets temporary directory for saving model.'''

  base_dirpath = 'workertemp_' + str(task_id)
  temp_dir = os.path.join(dirpath, base_dirpath)
  tf.io.gfile.makedirs(temp_dir)
  return temp_dir


def write_filepath(filepath, task_type, task_id):
  '''Helper function. Gets filepath to save model.'''

  dirpath = os.path.dirname(filepath)
  base = os.path.basename(filepath)
  if not _is_chief(task_type, task_id):
    dirpath = _get_temp_dir(dirpath, task_id)
  return os.path.join(dirpath, base)


def main():
  # Create strategy
  strategy = tf.distribute.MultiWorkerMirroredStrategy()

  # Get data
  global_batch_size = PER_REPLICA_BATCH_SIZE * strategy.num_replicas_in_sync
  train_data, number_of_classes = create_dataset(global_batch_size)

  # Wrap variable creation within strategy scope
  with strategy.scope():
    model = create_model(number_of_classes)

  model.fit(train_data, epochs=EPOCHS)

  # Determine type and task of the machine from
  # the strategy cluster resolver
  task_type, task_id = (strategy.cluster_resolver.task_type,
                        strategy.cluster_resolver.task_id)

  # Based on the type and task, write to the desired model path
  write_model_path = write_filepath(BUCKET, task_type, task_id)
  model.save(write_model_path)

if __name__ == "__main__":
    main()

Antes de compilar el contenedor, analicemos el código en detalle, que usa MultiWorkerMirroredStrategy de la API de tf.distribute.Strategy.

MultiWorkerMirroredStrategy está diseñado para distribuir el entrenamiento en varias máquinas de un clúster, cada una con varias GPUs. Si solo quieres distribuir el entrenamiento en varias GPUs en un solo dispositivo, consulta MirroredStrategy.

Hay algunos componentes en el código que son necesarios para que tu código funcione con MultiWorkerMirroredStrategy.

1. Los datos deben fragmentarse, lo que significa que a cada trabajador se le asigna un subconjunto de todo el conjunto de datos. Por lo tanto, en cada paso, cada trabajador procesará un tamaño de lote global de elementos del conjunto de datos que no se superponen. Esta fragmentación se hace automáticamente con tf.data.experimental.AutoShardPolicy, que se puede establecer en FILE o DATA. En este ejemplo, la función create_dataset() establece AutoShardPolicy en DATA porque el conjunto de datos de mandioca no se descarga como varios archivos. Sin embargo, si no configuraste la política en DATA, se activaría la política predeterminada de AUTO y el resultado final sería el mismo. Obtén más información sobre el fragmento de conjuntos de datos con MultiWorkerMirroredStrategy aquí.
2. En la función main(), se crea el objeto MultiWorkerMirroredStrategy. A continuación, agrega la creación de las variables de tu modelo al alcance de la estrategia. En este paso crucial, se indica a TensorFlow qué variables se deben duplicar en las réplicas.
3. El tamaño del lote se escala verticalmente según num_replicas_in_sync. Esto garantiza que cada réplica procese la misma cantidad de ejemplos en cada paso. Escalar el tamaño del lote es una práctica recomendada cuando se usan estrategias de paralelismo de datos síncronos en TensorFlow.
4. Guardar el modelo es un poco más complicado en el caso de varios trabajadores porque el destino necesita ser diferente para cada uno de ellos. El trabajador principal se guardará en el directorio de modelo deseado, mientras que los otros trabajadores guardarán el modelo en directorios temporales. Es importante que estos directorios temporales sean únicos para evitar que varios trabajadores escriban en la misma ubicación. La acción de guardar puede contener operaciones colectivas, lo que significa que todos los trabajadores deben realizarla, no solo el principal. Las funciones _is_chief(), _get_temp_dir(), write_filepath(), al igual que la función main() incluyen código estándar que ayuda a guardar el modelo.

Ten en cuenta que, si usaste MultiWorkerMirroredStrategy en un entorno diferente, es posible que hayas configurado la variable de entorno TF_CONFIG. Vertex AI configura TF_CONFIG automáticamente, por lo que no necesitas definir esta variable en cada máquina del clúster.

Paso 4: Compila el contenedor

Desde tu terminal, ejecuta lo siguiente para 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/multiworker:cassava"

Configura el Docker.

gcloud auth configure-docker

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

docker build ./ -t $IMAGE_URI

Por último, envía el contenedor a Google Container Registry:

docker push $IMAGE_URI

Ahora que enviaste el contenedor a Container Registry, puedes iniciar el trabajo de entrenamiento.

6. Ejecuta un trabajo de entrenamiento de varios trabajadores en Vertex AI

En este lab, se usa un entrenamiento personalizado mediante un contenedor personalizado en Google Container Registry, pero también puedes ejecutar un trabajo de entrenamiento con los contenedores creados previamente.

Para comenzar, ve a Entrenamiento (Training) de la sección de Vertex de la consola de Cloud.

Paso 1: Configura el trabajo de entrenamiento

Haz clic en Crear (Create) para ingresar los parámetros del trabajo de entrenamiento.

• En Conjunto de datos, selecciona No hay ningún conjunto de datos administrado.
• Selecciona Entrenamiento personalizado (avanzado) como método de entrenamiento y haz clic en Continuar.
• Ingresa multiworker-cassava (o el nombre que quieras asignarle a tu modelo) en Nombre del modelo.
• Haga clic en Continue.

En el paso Configuración del contenedor, selecciona Contenedor personalizado (Custom container):

En la primera casilla (Imagen de contenedor), ingresa el valor de la variable IMAGE_URI de la sección anterior. Debería ser gcr.io/your-cloud-project/multiworker:cassava, con el ID de tu proyecto. Deja el resto de los campos en blanco y haz clic en Continuar.

Para omitir el paso de hiperparámetros, vuelve a hacer clic en Continuar.

Paso 2: Configura el clúster de procesamiento

Vertex AI proporciona 4 grupos de trabajadores para abarcar los diferentes tipos de tareas de máquinas.

El grupo de trabajadores 0 configura el trabajador principal o programador. En MultiWorkerMirroredStrategy, todas las máquinas se designan como trabajadores, que son las máquinas físicas en las que se ejecuta el procesamiento replicado. Además de que cada máquina sea un trabajador, debe haber un trabajador que realice tareas adicionales, como guardar los puntos de control y escribir archivos de resumen en TensorBoard. Esta máquina se conoce como principal. Solo hay un trabajador principal, por lo que la cantidad de trabajadores en el grupo 0 siempre será 1.

En Procesamiento y precios, no cambies la región seleccionada y configura Grupo de trabajadores 0 de la siguiente manera:

Puedes optar por agregar GPU, pero ten en cuenta que, para MultiWorkerMirroredStrategy, cada máquina de tu clúster debería tener la misma cantidad de GPU. Si agregas GPUs, también aumentará el costo del trabajo de entrenamiento.

En el grupo de trabajadores 1, se configuran los trabajadores del clúster.

Configura el grupo de trabajadores 1 de la siguiente manera:

El clúster ahora está configurado para tener dos máquinas solo de CPU. Cuando se ejecute el código de la aplicación de entrenamiento, MultiWorkerMirroredStrategy distribuirá el entrenamiento en ambas máquinas.

MultiWorkerMirroredStrategy solo tiene los tipos de tareas de trabajador y trabajador principal, por lo que no es necesario configurar los grupos de trabajadores adicionales. Sin embargo, si usaras ParameterServerStrategy de TensorFlow, configurarías tus servidores de parámetros en el grupo de trabajadores 2. Además, si quisieras agregar un evaluador a tu clúster, configurarías esa máquina en el grupo de trabajadores 3.

Haz clic en Comenzar el entrenamiento para iniciar el trabajo de ajuste de hiperparámetros. En la sección Entrenamiento (Training) de la consola, en la pestaña CANALIZACIONES DE ENTRENAMIENTO (TRAINING PIPELINES), verás el trabajo iniciado recientemente:

El trabajo de entrenamiento tardará entre 30 y 35 minutos en completarse.

🎉 ¡Felicitaciones! 🎉

Aprendiste a usar Vertex AI para hacer lo siguiente:

• Iniciar un trabajo de entrenamiento de varios trabajadores para entrenar el código proporcionado en un contenedor personalizado. En este ejemplo, usaste un modelo de TensorFlow, pero puedes entrenar un modelo creado con cualquier framework a través de contenedores personalizados o integrados.

Para obtener más información sobre las distintas partes de Vertex, consulte la documentación.

7. Use el SDK de Vertex (opcional)

En la sección anterior, se mostró cómo iniciar el trabajo de entrenamiento con la IU. En esta sección, verás una forma alternativa de enviar el trabajo de entrenamiento mediante la API de Vertex Python.

Regresa a tu instancia de notebook y crea un notebook de TensorFlow 2 desde Selector (Launcher):

Importa el SDK de Vertex AI.

from google.cloud import aiplatform

Para iniciar el trabajo de entrenamiento de varios trabajadores, primero debes definir la especificación del grupo de trabajadores. Ten en cuenta que el uso de GPUs en la especificación es completamente opcional y puedes quitar accelerator_type y accelerator_count si deseas un clúster solo de CPU, como se muestra en la sección anterior.

# The spec of the worker pools including machine type and Docker image
# Be sure to replace {YOUR-PROJECT-ID} with your project ID.
worker_pool_specs=[
     {
        "replica_count": 1,
        "machine_spec": {
          "machine_type": "n1-standard-8", "accelerator_type": "NVIDIA_TESLA_V100", "accelerator_count": 1
        },
        "container_spec": {"image_uri": "gcr.io/{YOUR-PROJECT-ID}/multiworker:cassava"}
      },
      {
        "replica_count": 1,
        "machine_spec": {
          "machine_type": "n1-standard-8", "accelerator_type": "NVIDIA_TESLA_V100", "accelerator_count": 1
        },
        "container_spec": {"image_uri": "gcr.io/{YOUR-PROJECT-ID}/multiworker:cassava"}
      }
]

A continuación, crea y ejecuta un CustomJob. Deberás reemplazar {YOUR_BUCKET} por un bucket de tu proyecto para la etapa de pruebas. Puedes usar el mismo bucket que creaste antes.

# Replace YOUR_BUCKET
my_multiworker_job = aiplatform.CustomJob(display_name='multiworker-cassava-sdk',
                              worker_pool_specs=worker_pool_specs,
                              staging_bucket='gs://{YOUR_BUCKET}')

my_multiworker_job.run()

En la sección Entrenamiento de la consola, en la pestaña TRABAJOS PERSONALIZADOS, verás el trabajo de entrenamiento:

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 la consola de Cloud, navega a Almacenamiento, selecciona tu bucket y haz clic en Borrar (Delete):