Vertex AI: Usa el empaquetado automático para ajustar BERT con Hugging Face en Vertex AI Training

1. Descripción general

En este lab, aprenderás a ejecutar un trabajo de entrenamiento personalizado en Vertex AI Training con la función de empaquetado automático. Los trabajos de entrenamiento personalizados en Vertex AI usan contenedores. Si no quieres compilar tu propia imagen, puedes usar auotpackaging, el cual compilará una imagen personalizada de Docker basada en tu código, la enviará a Container Registry y, luego, iniciará un CustomJob basado en la imagen.

Qué aprenderá

Aprenderás a hacer lo siguiente:

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

2. Descripción general del caso de uso

Con las bibliotecas de Hugging Face, podrás ajustar un modelo BERT en el conjunto de datos de IMDB. El modelo predecirá si la opinión sobre una película es positiva o negativa. El conjunto de datos se descargará de la biblioteca de conjuntos de datos de Hugging Face y el modelo BERT de la biblioteca de transformadores de Hugging Face.

3. 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 Training y Workbench.

Descripción general del producto Vertex

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 Vertex AI

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

Panel de Vertex AI

Paso 3: Habilita la API de Container Registry

Navega a Container Registry y selecciona Habilitar si aún no lo has hecho. La usarás para crear un contenedor para tu trabajo de entrenamiento personalizado.

Paso 4: Crea una instancia de Vertex AI Workbench

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

Menú Vertex AI

Allí, haz clic en NOTEBOOKS ADMINISTRADOS (MANAGED NOTEBOOKS):

Notebooks_UI

Luego, selecciona NUEVO NOTEBOOK (NEW NOTEBOOK).

new_notebook

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

create_notebook

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.

idle_timeout

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

Luego, haz clic en Crear.

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

open_jupyterlab

La primera vez que uses una instancia nueva, se te solicitará que te autentiques.

autenticar

5. Escribe el código de entrenamiento

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

launcher_terminal

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

mkdir autopkg-codelab
cd autopkg-codelab

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 autopkg-codelab/:

+ 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.

import argparse

import tensorflow as tf
from datasets import load_dataset
from transformers import AutoTokenizer
from transformers import TFAutoModelForSequenceClassification

CHECKPOINT = "bert-base-cased"

def get_args():
  '''Parses args.'''

  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--epochs',
      required=False,
      default=3,
      type=int,
      help='number of epochs')
  parser.add_argument(
      '--job_dir',
      required=True,
      type=str,
      help='bucket to store saved model, include gs://')
  args = parser.parse_args()
  return args


def create_datasets():
    '''Creates a tf.data.Dataset for train and evaluation.'''

    raw_datasets = load_dataset('imdb')
    tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
    tokenized_datasets = raw_datasets.map((lambda examples: tokenize_function(examples, tokenizer)), batched=True)

    # To speed up training, we use only a portion of the data.
    # Use full_train_dataset and full_eval_dataset if you want to train on all the data.
    small_train_dataset = tokenized_datasets['train'].shuffle(seed=42).select(range(1000))
    small_eval_dataset = tokenized_datasets['test'].shuffle(seed=42).select(range(1000))
    full_train_dataset = tokenized_datasets['train']
    full_eval_dataset = tokenized_datasets['test']

    tf_train_dataset = small_train_dataset.remove_columns(['text']).with_format("tensorflow")
    tf_eval_dataset = small_eval_dataset.remove_columns(['text']).with_format("tensorflow")

    train_features = {x: tf_train_dataset[x] for x in tokenizer.model_input_names}
    train_tf_dataset = tf.data.Dataset.from_tensor_slices((train_features, tf_train_dataset["label"]))
    train_tf_dataset = train_tf_dataset.shuffle(len(tf_train_dataset)).batch(8)

    eval_features = {x: tf_eval_dataset[x] for x in tokenizer.model_input_names}
    eval_tf_dataset = tf.data.Dataset.from_tensor_slices((eval_features, tf_eval_dataset["label"]))
    eval_tf_dataset = eval_tf_dataset.batch(8)

    return train_tf_dataset, eval_tf_dataset


def tokenize_function(examples, tokenizer):
    '''Tokenizes text examples.'''

    return tokenizer(examples['text'], padding='max_length', truncation=True)


def main():
    args = get_args()
    train_tf_dataset, eval_tf_dataset = create_datasets()
    model = TFAutoModelForSequenceClassification.from_pretrained(CHECKPOINT, num_labels=2)

    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=0.01),
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=tf.metrics.SparseCategoricalAccuracy(),
    )

    model.fit(train_tf_dataset, validation_data=eval_tf_dataset, epochs=args.epochs)
    model.save(f'{args.job_dir}/model_output')


if __name__ == "__main__":
    main()

Algunos aspectos que debes tener en cuenta sobre el código:

  • CHECKPOINT es el modelo que queremos ajustar. En este caso, usamos BERT.
  • El método TFAutoModelForSequenceClassification cargará la arquitectura del modelo de lenguaje especificado y los pesos en TensorFlow, y agregará un encabezado de clasificación en la parte superior con pesos inicializados de forma aleatoria. En este caso, tenemos un problema de clasificación binaria (positivo o negativo), por lo que especificamos num_labels=2 para este clasificador.

6. Crear contenedores y ejecutar código de entrenamiento de manera local

Puedes usar el comando gcloud ai custom-jobs local-run para compilar una imagen de contenedor de Docker basada en el código de entrenamiento y ejecutarla como un contenedor en tu máquina local. La ejecución local de un contenedor ejecuta tu código de entrenamiento de manera similar a la que se ejecuta en Vertex AI Training y puede ayudarte a depurar problemas con tu código antes de realizar un entrenamiento personalizado en Vertex AI.

En el trabajo de entrenamiento, exportaremos el modelo 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:

PROJECT_ID='your-cloud-project'

Luego, crea un bucket. Si ya tienes un bucket, puedes usarlo en su lugar.

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

Cuando ejecutemos el trabajo de entrenamiento personalizado en Vertex AI Training, usaremos una GPU. Sin embargo, como no especificamos nuestra instancia de Workbench con GPU, usaremos una imagen basada en CPU para las pruebas locales. En este ejemplo, usamos un contenedor creado previamente de Vertex AI Training.

Ejecuta el siguiente comando para establecer el URI de una imagen de Docker que se usará como base del contenedor.

BASE_CPU_IMAGE=us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-7:latest

Luego, establece un nombre para la imagen de Docker resultante compilada por el comando de ejecución local.

OUTPUT_IMAGE=$PROJECT_ID-local-package-cpu:latest

Nuestro código de entrenamiento usa los conjuntos de datos y las bibliotecas de transformadores de Hugging Face. Estas bibliotecas no se incluyen en la imagen que seleccionamos como imagen base, por lo que tendremos que proporcionarlas como requisitos. Para ello, crearemos un archivo requirements.txt en nuestro directorio autopkg-codelab.

Asegúrate de estar en el directorio autopkg-codelab y escribe lo siguiente en la terminal.

touch requirements.txt

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

+ requirements.txt
+ trainer/
    + task.py

Abre el archivo de requisitos y pega el siguiente

datasets==1.18.2
transformers==4.16.2

Por último, ejecuta el comando gcloud ai custom-jobs local-run para iniciar el entrenamiento en nuestra instancia administrada de Workbench.

gcloud ai custom-jobs local-run \
--executor-image-uri=$BASE_CPU_IMAGE \
--python-module=trainer.task \
--output-image-uri=$OUTPUT_IMAGE \
-- \
--job_dir=$BUCKET_NAME

Deberías ver la imagen de Docker que se está compilando. Las dependencias que agregamos al archivo requirements.txt se instalarán mediante pip. Esto puede tardar unos minutos en completarse la primera vez que ejecutes este comando. Una vez que se compile la imagen, comenzará a ejecutarse el archivo task.py y verás el entrenamiento del modelo. Debería ver algo como esto:

local_training

Debido a que no estamos usando una GPU localmente, el entrenamiento del modelo llevará mucho tiempo. Puede presionar Ctrl + C y cancelar el entrenamiento local en lugar de esperar a que se complete el trabajo.

Ten en cuenta que si deseas realizar más pruebas, también puedes ejecutar directamente la imagen compilada anteriormente, sin reempaquetar.

gcloud beta ai custom-jobs local-run \
--executor-image-uri=$OUTPUT_IMAGE \
-- \
--job_dir=$BUCKET_NAME \
--epochs=1

7. Crea un trabajo personalizado

Ahora que probamos el modo local, usaremos la función de empaquetado automático para iniciar nuestro trabajo de entrenamiento personalizado en Vertex AI Training. Con un solo comando, esta función hará lo siguiente:

  • Compilar una imagen de Docker personalizada basada en tu código
  • Envía la imagen a Container Registry.
  • Iniciar una CustomJob basada en la imagen

Regresa a la terminal y sube un nivel por encima de tu directorio autopkg-codelab.

+ autopkg-codelab
  + requirements.txt
  + trainer/
      + task.py

Especifica la imagen de GPU de TensorFlow precompilada con Vertex AI Training como imagen base para el trabajo de entrenamiento personalizado.

BASE_GPU_IMAGE=us-docker.pkg.dev/vertex-ai/training/tf-gpu.2-7:latest

Luego, ejecuta el comando gcloud ai custom-jobs create. Primero, este comando compilará una imagen de Docker personalizada basada en el código de entrenamiento. La imagen base es el contenedor precompilado de Vertex AI Training que establecimos como BASE_GPU_IMAGE. Luego, la función de empaquetado automático instalará por pip las bibliotecas de conjuntos de datos y transformadores como se especifica en nuestro archivo requirements.txt.

gcloud ai custom-jobs create \
--region=us-central1 \
--display-name=fine_tune_bert \
--args=--job_dir=$BUCKET_NAME \
--worker-pool-spec=machine-type=n1-standard-4,replica-count=1,accelerator-type=NVIDIA_TESLA_V100,executor-image-uri=$BASE_GPU_IMAGE,local-package-path=autopkg-codelab,python-module=trainer.task

Veamos el argumento worker-pool-spec. Esto define la configuración del grupo de trabajadores que usa el trabajo personalizado. Puedes especificar varias especificaciones de grupos de trabajadores a fin de crear un trabajo personalizado con varios grupos de trabajadores para el entrenamiento distribuido. En este ejemplo, solo especificamos un único grupo de trabajadores, ya que nuestro código de entrenamiento no está configurado para entrenamiento distribuido.

Estos son algunos de los campos clave de esta especificación:

  • machine-type (obligatorio): Es el tipo de máquina. Haz clic aquí para conocer los tipos compatibles.
  • replica-count: Es la cantidad de réplicas de trabajadores que se deben usar para este grupo de trabajadores. De forma predeterminada, el valor es 1.
  • accelerator-type: Es el tipo de GPU. Haz clic aquí para conocer los tipos compatibles. En este ejemplo, especificamos una GPU NVIDIA Tesla V100.
  • accelerator-count: Es la cantidad de GPU que se usará para cada VM en el grupo de trabajadores. El valor es 1 de forma predeterminada.
  • executor-image-uri: Es el URI de una imagen de contenedor que ejecutará el paquete proporcionado. Esto se establece en nuestra imagen base.
  • local-package-path: Es la ruta de acceso local de una carpeta que contiene el código de entrenamiento.
  • python-module: Es el nombre del módulo de Python que se ejecutará dentro del paquete proporcionado.

Al igual que cuando ejecutaste el comando local, verás que se compila la imagen de Docker y, luego, se inicia el trabajo de entrenamiento. Excepto que no veas el resultado del trabajo de entrenamiento, verás el siguiente mensaje, que confirma que se inició el trabajo de entrenamiento. Ten en cuenta que la primera vez que ejecutes el comando custom-jobs create, la compilación y el envío de la imagen pueden tardar unos minutos.

training_started

Regresa a la sección Vertex AI Training de la consola de Cloud y, en TRABAJOS PERSONALIZADOS, deberías ver tu trabajo en ejecución.

training_job

El trabajo tardará unos 20 minutos en completarse.

Una vez que se complete el proceso, deberías ver los siguientes artefactos del modelo guardado en el directorio model_output de tu bucket.

model_output

🎉 ¡Felicitaciones! 🎉

Aprendiste a usar Vertex AI para hacer lo siguiente:

  • Crear contenedores y ejecutar código de entrenamiento de manera local
  • Envía trabajos de entrenamiento a Vertex AI Training mediante empaquetado automático

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).

borrar

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):

Borrar almacenamiento