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

1. Resumen

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

Qué aprenderá

Obtendrás información sobre cómo hacer las siguientes acciones:

• Modifica el código de la aplicación de entrenamiento para el entrenamiento de varios trabajadores
• Configurar e iniciar un trabajo de entrenamiento de varios trabajadores desde la IU de Vertex AI
• Configure e inicie 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 $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 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 los productos destacados a continuación: Capacitación y Workbench

3. Descripción general de caso de uso

En este lab, utilizará el aprendizaje por transferencia para entrenar un modelo de clasificación de imágenes en el conjunto de datos de yuca, 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 ImageNet.

¿Por qué debería realizar el entrenamiento distribuido?

Si tienes una sola GPU, TensorFlow usará este acelerador para acelerar el entrenamiento de modelos sin que tengas que realizar ninguna acción adicional. Sin embargo, si deseas obtener un aumento adicional del uso de varias GPU en una sola máquina o en varias máquinas (cada una con posiblemente varias GPU), deberás usar tf.distribute, que es la biblioteca de TensorFlow para ejecutar cálculos. varios dispositivos. En una máquina en la que TensorFlow puede ejecutar operaciones, un dispositivo hace referencia a una CPU o un acelerador, como GPU o TPU.

La forma más sencilla de comenzar el entrenamiento distribuido es usar una sola máquina con varios dispositivos de GPU. Una estrategia de distribución de TensorFlow desde el módulo tf.distribute administrará la coordinación de la distribución de datos y las actualizaciones de gradientes en todas las GPU. Si ya domina el entrenamiento de host único y desea escalar aún más, agregar varias máquinas a su clúster puede ayudarlo a mejorar aún más el rendimiento. Puedes usar un clúster de máquinas que son solo de CPU, o que cada una tiene una o más GPU. En este lab, se trata este ú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íncrono que puedes usar con solo algunos cambios de 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 de trabajador calcula los datos hacia adelante y hacia atrás a través del 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 "reduce". Luego, el optimizador realiza las actualizaciones de parámetros con estos gradientes reducidos y, de esta manera, mantiene los dispositivos sincronizados. Para obtener más información sobre el entrenamiento distribuido con TensorFlow, mira el siguiente video:

4. Configure el entorno

Para ejecutar este codelab, necesitarás un proyecto de Google Cloud Platform con facturación habilitada. 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á habilitado). La necesitarás para crear tu instancia de notebook.

Paso 2: Habilita la API de Container Registry

Navegue a Container Registry y seleccione Habilitar si aún no lo ha hecho. Usarás esto a fin de crear un contenedor para tu trabajo de entrenamiento personalizado.

Paso 3: Habilita la API de Vertex AI

Navegue hasta la sección de Vertex AI en Cloud Console y haga clic en Habilitar API de Vertex AI.

Paso 4: Crea una instancia de Vertex AI Workbench

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

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

Una vez habilitada, haga clic en NOTED NOTEBOOK:

Luego, selecciona NUEVO NOTEBOOK.

Asigne un nombre a su notebook y, luego, haga clic en Configuración avanzada.

En Configuración avanzada, habilita el cierre inactivo y establece la cantidad de minutos en 60. Esto significa que el notebook se cerrará automáticamente cuando no esté en uso para que no se generen costos innecesarios.

En Seguridad, selecciona "Habilitar terminal" si aún no está habilitado.

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

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

Una vez que se haya creado la instancia, selecciona Abrir JupyterLab.

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

5. Crea contenedores para el código de aplicaciones de entrenamiento

Para enviar este trabajo de entrenamiento a Vertex, coloque el código de su aplicación de entrenamiento en un contenedor de Docker y lo enviará a Google Container Registry. Con este enfoque, puedes entrenar un modelo creado con cualquier framework.

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

Crea un directorio nuevo llamado cassava y cd en él:

mkdir cassava
cd cassava

Paso 1: Cree un Dockerfile

El primer paso para crear un contenedor de código es crear un Dockerfile. En el Dockerfile, incluirá 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.

Desde tu terminal, crea un Dockerfile vacío:

touch Dockerfile

Abra el Dockerfile y copie 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 de Enterprise 2.7 de Deep Learning Container para TensorFlow. Los contenedores de aprendizaje profundo de Google Cloud tienen muchos frameworks comunes de AA y ciencia de datos preinstalados. Después de descargar esa imagen, este Dockerfile configura el punto de entrada para el código de entrenamiento. Aún no has creado 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á el modelo entrenado de TensorFlow a un bucket de Cloud Storage. 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 por 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'

A continuación, ejecute el siguiente comando en su terminal para crear un nuevo bucket en su proyecto.

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

Paso 3: Agrega el código de entrenamiento de modelos

Desde su terminal, ejecute el siguiente comando a fin de crear un directorio para el código de entrenamiento y un archivo de Python en el que agregará 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 siguiente código. 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 en detalle el código, 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 posibles GPU múltiples. Si solo deseas distribuir el entrenamiento entre varias GPU 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 del lote global de elementos del conjunto de datos que no se superpongan. Esta fragmentación ocurre automáticamente con tf.data.experimental.AutoShardPolicy, que se puede configurar como FILE o DATA. En este ejemplo, la función create_dataset() establece AutoShardPolicy en DATA porque el conjunto de datos de la yuca no se descarga como varios archivos. Sin embargo, si no estableces la política como DATA, se activará la política predeterminada AUTO, y el resultado final será el mismo. Puedes obtener más información sobre la fragmentación del conjunto de datos con MultiWorkerMirroredStrategy aquí.
2. En la función main(), se crea el objeto MultiWorkerMirroredStrategy. A continuación, une la creación de sus variables de modelo dentro del alcance de la estrategia. Este paso crucial le indica a TensorFlow qué variables deben duplicarse en las réplicas.
3. El tamaño del lote se escala verticalmente según el num_replicas_in_sync. Esto garantiza que cada réplica procese la misma cantidad de ejemplos en cada paso. El escalamiento del tamaño del lote es una práctica recomendada cuando se usan estrategias de paralelismo de datos síncrono en TensorFlow.
4. Guardar el modelo es un poco más complicado en el caso de varios trabajadores porque el destino debe ser diferente para cada uno de los trabajadores. El trabajador principal guardará el directorio en el directorio del modelo deseado, mientras que los demás trabajadores guardarán el modelo en los directorios temporales. Es importante que estos directorios temporales sean únicos para evitar que varios trabajadores escriban en la misma ubicación. Guardar puede contener operaciones colectivas, lo que significa que todos los trabajadores deben guardar y no solo el principal. Las funciones _is_chief(), _get_temp_dir() y write_filepath(), así como 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 de manera automática, 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 a fin de definir una variable de entorno para tu proyecto y asegúrate de reemplazar your-cloud-project por 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'

Defina una variable con el URI de su imagen de contenedor en Google Container Registry:

IMAGE_URI="gcr.io/$PROJECT_ID/multiworker:cassava"

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íela a Google Container Registry:

docker push $IMAGE_URI

Con el contenedor enviado a Container Registry, ya está listo para iniciar un trabajo de entrenamiento.

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

Este lab utiliza el entrenamiento personalizado a través de un contenedor personalizado en Google Container Registry, pero también puede ejecutar un trabajo de entrenamiento con los contenedores previamente compilados.

Para comenzar, navega a la sección Entrenamiento en la sección Vertex de Cloud Console:

Paso 1: Configura el trabajo de entrenamiento

Haz clic en Crear para ingresar los parámetros de tu trabajo de entrenamiento.

• En Conjunto de datos, selecciona Sin conjunto de datos administrado.
• Luego, selecciona Entrenamiento personalizado (avanzado) como método de entrenamiento y haz clic en Continuar.
• Ingresa multiworker-cassava (o el nombre que quieras llamar al modelo) en Model name.
• Haga clic en Continue.

En el paso de Configuración del contenedor, seleccione Contenedor personalizado:

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

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

Paso 2: Configura el clúster de procesamiento

Vertex AI ofrece 4 grupos de trabajadores para abarcar los diferentes tipos de tareas de máquina.

En el grupo de trabajadores 0, se configura el primario, el jefe, el programador o el "principal". En MultiWorkerMirroredStrategy, todas las máquinas se designan como trabajadores, que son las máquinas físicas en las que se ejecuta el cálculo replicado. Además de que cada máquina sea un trabajador, debe haber un trabajador que asuma trabajo adicional, como guardar puntos de control y escribir archivos de resumen en TensorBoard. Esta máquina se conoce como la principal. Solo hay un trabajador principal, por lo que el recuento de trabajadores del grupo de trabajadores 0 siempre será 1.

En Compute and pricing, deje la región seleccionada tal como está y configure el 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 debe tener la misma cantidad de GPU. Agregar GPU también aumentará el costo del trabajo de entrenamiento.

El grupo de trabajadores 1 es donde configuras los trabajadores para tu clúster.

Configura el Grupo de trabajadores 1 de la siguiente manera:

El clúster ahora está configurado para tener dos máquinas solo con 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 principales y de trabajador, por lo que no es necesario configurar los grupos de trabajadores adicionales. Sin embargo, si usaras el ParameterServerStrategy de TensorFlow, configurarías tus servidores de parámetros en el grupo de trabajadores 2. Si desea agregar un evaluador a su clúster, debe configurar esa máquina en el grupo de trabajadores 3.

Haz clic en Iniciar entrenamiento para iniciar el trabajo de ajuste de hiperparámetros. En la sección Training de su consola, en la pestaña TRAINING PIPELINES, verá su trabajo recién iniciado:

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 el código de entrenamiento proporcionado en un contenedor personalizado En este ejemplo, usó un modelo de TensorFlow, pero puede entrenar un modelo creado con cualquier framework mediante contenedores personalizados o integrados.

Si quieres obtener más información sobre las distintas partes de Vertex, consulta la documentación.

7. Usa el SDK de Vertex (opcional)

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

Regresa a tu instancia de notebook y crea un notebook de TensorFlow 2 desde el Launcher:

Importa el SDK de Vertex AI.

from google.cloud import aiplatform

Para iniciar el trabajo de entrenamiento de varios trabajadores, primero debe definir la especificación del grupo de trabajadores. Ten en cuenta que el uso de GPU en la especificación es completamente opcional y puedes quitar accelerator_type y accelerator_count si deseas un clúster de solo 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 objeto CustomJob. Deberás reemplazar {YOUR_BUCKET} por un bucket de tu proyecto para la etapa de pruebas. Puedes usar el mismo bucket que creaste anteriormente.

# 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 su consola, en la pestaña CUSTOM JOBS, verá su trabajo de entrenamiento:

8. 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 desea cerrar la instancia de forma manual, haga clic en el botón Detener en la sección Vertex AI Workbench de la consola. Si deseas borrar el notebook por completo, haz clic en el botón Borrar.

Para borrar el bucket de Storage, en el menú de navegación de Cloud Console, navegue a Storage, seleccione su bucket y haga clic en Borrar: