Vertex AI: use o empacotamento automático para ajustar o Bert com o Hugging Face no Vertex AI Training

1. Visão geral

Neste laboratório, você vai aprender a executar um job de treinamento personalizado no Vertex AI Training com o recurso de empacotamento automático. Os jobs de treinamento personalizado na Vertex AI usam contêineres. Se você não quiser criar sua própria imagem, use o empacotamento automático, que vai criar uma imagem Docker personalizada com base no seu código, enviar a imagem para o Container Registry e iniciar um CustomJob com base na imagem.

Conteúdo do laboratório

Você vai aprender a:

O custo total da execução deste laboratório no Google Cloud é de aproximadamente US$2.

2. Visão geral do caso de uso

Usando bibliotecas do Hugging Face, você vai ajustar um modelo Bert no conjunto de dados do IMDB. O modelo vai prever se uma crítica de filme é positiva ou negativa. O conjunto de dados será baixado da biblioteca de conjuntos de dados do Hugging Face e o modelo Bert da biblioteca de transformadores do Hugging Face.

3. Introdução à Vertex AI

Este laboratório usa a mais nova oferta de produtos de IA disponível no Google Cloud. A Vertex AI integra as ofertas de ML do Google Cloud em uma experiência de desenvolvimento intuitiva. Anteriormente, modelos treinados com o AutoML e modelos personalizados eram acessíveis por serviços separados. A nova oferta combina ambos em uma única API, com outros novos produtos. Você também pode migrar projetos existentes para a Vertex AI. Se você tiver algum feedback, consulte a página de suporte.

A Vertex AI tem vários produtos diferentes incluídos para dar suporte a fluxos de trabalho integrais de ML. O foco deste laboratório é o Treinamento e o Workbench.

Visão geral do produto Vertex

4. Configurar o ambiente

Para executar este codelab, você vai precisar de um projeto do Google Cloud Platform com o faturamento ativado. Para criar um projeto, siga estas instruções.

Etapa 1: ativar a API Compute Engine

Acesse o Compute Engine e selecione Ativar, caso essa opção ainda não esteja ativada.

Etapa 2: ativar a API Vertex AI

Navegue até a seção "Vertex AI" do Console do Cloud e clique em Ativar API Vertex AI.

Painel da Vertex AI

Etapa 3: ativar a API Container Registry

Navegue até o Container Registry e selecione Ativar se ele ainda não estiver. Use isso para criar um contêiner para seu job de treinamento personalizado.

Etapa 4: criar uma instância do Vertex AI Workbench

Na seção Vertex AI do Console do Cloud, clique em "Workbench":

Menu da Vertex AI

Em seguida, clique em NOTEBOOKS GERENCIADOS:

Notebooks_UI

Em seguida, selecione NOVO NOTEBOOK.

new_notebook

Dê um nome ao notebook e clique em Configurações avançadas.

create_notebook

Em "Configurações avançadas", ative o encerramento inativo e defina o número de minutos como 60. Isso significa que o notebook será desligado automaticamente quando não estiver sendo usado.

idle_timeout

Você pode manter as outras configurações avançadas como estão.

Em seguida, clique em Criar.

Quando a instância tiver sido criada, selecione Abrir o JupyterLab.

open_jupyterlab

Na primeira vez que usar uma nova instância, você vai receber uma solicitação de autenticação.

autenticar

5. Escrever o código de treinamento

Para começar, no menu de acesso rápido, abra uma janela de terminal na instância do notebook:

launcher_terminal

Crie um novo diretório chamado autopkg-codelab e coloque cd nele.

mkdir autopkg-codelab
cd autopkg-codelab

No seu Terminal, execute o seguinte para criar um diretório para o código de treinamento e um arquivo Python onde você vai adicionar o código:

mkdir trainer
touch trainer/task.py

Agora você deve ter o seguinte no diretório autopkg-codelab/:

+ trainer/
    + task.py

Depois abra o arquivo task.py que você acabou de criar e copie o código abaixo.

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

Algumas observações sobre o código:

  • CHECKPOINT é o modelo que queremos ajustar. Nesse caso, usamos o Bert.
  • O método TFAutoModelForSequenceClassification vai carregar a arquitetura do modelo de linguagem especificada + pesos no TensorFlow e adicionar um cabeçalho de classificação na parte de cima com pesos inicializados aleatoriamente. Nesse caso, temos um problema de classificação binária (positivo ou negativo), então especificamos num_labels=2 para esse classificador.

6. Coloque o código de treinamento em contêineres e faça a execução localmente

É possível usar o comando gcloud ai custom-jobs local-run para criar uma imagem de contêiner do Docker com base no código de treinamento e executar a imagem como um contêiner na máquina local. A execução local de um contêiner executa o código de treinamento de maneira semelhante à execução no Vertex AI Training e pode ajudar a depurar problemas com o código antes de executar o treinamento personalizado na Vertex AI.

No job de treinamento, vamos exportar o modelo treinado para um bucket do Cloud Storage. No Terminal, execute o comando a seguir e defina uma variável env para o projeto. Lembre-se de substituir your-cloud-project pelo ID do projeto.

PROJECT_ID='your-cloud-project'

Em seguida, crie um bucket. Se você já tiver um bucket, use-o.

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

Quando executarmos o job de treinamento personalizado no Vertex AI Training, vamos usar uma GPU. No entanto, como não especificamos a instância do Workbench com GPUs, vamos usar uma imagem baseada em CPU para testes locais. Neste exemplo, usamos um contêiner pré-criado do Vertex AI Training.

Execute o comando a seguir para definir o URI de uma imagem Docker a ser usada como base do contêiner.

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

Em seguida, defina um nome para a imagem Docker resultante criada pelo comando de execução local.

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

Nosso código de treinamento usa os conjuntos de dados do Hugging Face e as bibliotecas de transformadores. Essas bibliotecas não estão incluídas na imagem que selecionamos como imagem de base. Portanto, precisamos fornecê-las como requisitos. Para fazer isso, vamos criar um arquivo requirements.txt no diretório autopkg-codelab.

Verifique se você está no diretório autopkg-codelab e digite o seguinte no terminal.

touch requirements.txt

Agora você deve ter o seguinte no diretório autopkg-codelab:

+ requirements.txt
+ trainer/
    + task.py

Abra o arquivo de requisitos e cole o seguinte

datasets==1.18.2
transformers==4.16.2

Por fim, execute o comando gcloud ai custom-jobs local-run para iniciar o treinamento na instância gerenciada do 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

Você verá a imagem Docker sendo criada. As dependências que adicionamos ao arquivo requirements.txt serão instaladas pelo pip. Isso pode levar alguns minutos para ser concluído na primeira vez que você executar esse comando. Depois que a imagem for criada, o arquivo task.py vai começar a ser executado e o treinamento de modelo será exibido. Você vai conferir algo parecido com:

local_training

Como não estamos usando uma GPU localmente, o treinamento de modelo vai levar muito tempo. Você pode pressionar Ctrl + c e cancelar o treinamento local em vez de esperar que o job seja concluído.

Se você quiser fazer mais testes, também poderá executar diretamente a imagem criada acima, sem reempacotar.

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

7. Crie um job personalizado

Agora que testamos o modo local, vamos usar o recurso de empacotamento automático para iniciar nosso job de treinamento personalizado no Vertex AI Training. Com um único comando, esse recurso vai:

  • Crie uma imagem Docker personalizada com base no seu código.
  • Envie a imagem para o Container Registry
  • Inicia um CustomJob com base na imagem.

Volte ao terminal e cd um nível acima do diretório autopkg-codelab.

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

Especifique a imagem de GPU do TensorFlow pré-criada do Vertex AI Training como a imagem de base para o job de treinamento personalizado.

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

Em seguida, execute o comando gcloud ai custom-jobs create. Primeiro, esse comando vai criar uma imagem Docker personalizada com base no código de treinamento. A imagem de base é o contêiner pré-criado do Vertex AI Training que definimos como BASE_GPU_IMAGE. O recurso de empacotamento automático vai instalar os conjuntos de dados e as bibliotecas de transformadores conforme especificado no arquivo 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

Vamos analisar o argumento worker-pool-spec. Ele define a configuração do pool de workers usada pelo job personalizado. É possível especificar várias especificações de pool de workers para criar um job personalizado com vários pools de workers para treinamento distribuído. Neste exemplo, especificamos apenas um pool de workers, já que nosso código de treinamento não está configurado para treinamento distribuído.

Confira alguns dos principais campos dessa especificação:

  • machine-type (obrigatório): o tipo da máquina. Clique aqui para conferir os tipos compatíveis.
  • replica-count: o número de réplicas de worker a serem usadas para esse pool de workers. Por padrão, o valor é 1.
  • accelerator-type: o tipo de GPUs. Clique aqui para conferir os tipos compatíveis. Neste exemplo, especificamos uma GPU NVIDIA Tesla V100.
  • accelerator-count: o número de GPUs para cada VM no pool de workers a ser usado. Por padrão, o valor é 1.
  • executor-image-uri: o URI de uma imagem de contêiner que vai executar o pacote fornecido. Ele está definido como nossa imagem de base.
  • local-package-path: o caminho local de uma pasta que contém o código de treinamento.
  • python-module: o nome do módulo Python a ser executado no pacote fornecido.

Assim como quando você executou o comando local, a imagem Docker será criada e o job de treinamento será iniciado. No entanto, em vez de conferir a saída do job de treinamento, você vai conferir a mensagem a seguir confirmando que o job de treinamento foi iniciado. Na primeira vez que você executar o comando custom-jobs create, pode levar alguns minutos para que a imagem seja criada e enviada.

training_started

Volte para a seção "Vertex AI Training" do console do Cloud e, em JOBS PERSONALIZADOS , você vai conferir o job em execução.

training_job

O job vai levar cerca de 20 minutos para ser concluído.

Quando concluído, você vai conferir os artefatos do modelo salvo no diretório model_output no bucket.

model_output

Parabéns! 🎉

Você aprendeu a usar a Vertex AI para:

  • Coloque o código de treinamento em contêineres e faça a execução localmente
  • Enviar jobs de treinamento para o Vertex AI Training com empacotamento automático

Para saber mais sobre partes diferentes da Vertex AI, acesse a documentação.

8. Limpeza

Como configuramos o notebook para expirar após 60 minutos de inatividade, não precisamos nos preocupar em desligar a instância. Para encerrar a instância manualmente, clique no botão "Parar" na seção "Vertex AI Workbench" do console. Se quiser excluir o notebook completamente, clique no botão "Excluir".

excluir

Para excluir o bucket do Storage, use o menu de navegação do console do Cloud, acesse o Storage, selecione o bucket e clique em "Excluir":

Excluir armazenamento