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 auotpackaging, 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ê ajustará um modelo Bert no conjunto de dados do IMDB (links em inglês). O modelo prevê se uma crítica de filme é positiva ou negativa. O download do conjunto de dados será feito na biblioteca de conjuntos de dados Hugging Face e no modelo Bert da biblioteca de transformadores 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. 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

Depois, 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. Neste caso, usamos o Bert.
  • O método TFAutoModelForSequenceClassification vai carregar a arquitetura e os pesos do modelo de linguagem especificados no TensorFlow e adicionar um cabeçalho de classificação com pesos inicializados aleatoriamente. Neste caso, temos um problema de classificação binária (positivo ou negativo) e, por isso, especificamos num_labels=2 para esse classificador.

6. Conteinerizar e executar o código de treinamento 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. Além disso, ela ajuda a depurar problemas no 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, fique à vontade para usá-lo.

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

Quando executarmos o job de treinamento personalizado no Vertex AI Training, usaremos uma GPU. No entanto, como não especificamos nossa instância do Workbench com GPUs, usaremos 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 do Docker que 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 do 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 e as bibliotecas de transformadores do Hugging Face. Essas bibliotecas não estão incluídas na imagem selecionada como nossa imagem de base, portanto, precisaremos fornecê-las como requisitos. Para fazer isso, vamos criar um arquivo requirements.txt no nosso 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 este código

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 do Docker sendo criada. As dependências que adicionamos ao arquivo requirements.txt terão o pip instalado. Isso pode levar alguns minutos para ser concluído na primeira vez que você executar este comando. Depois que a imagem for criada, o arquivo task.py começará a ser executado e você verá o treinamento do modelo. Você verá algo como:

local_training

Como não estamos usando uma GPU localmente, o treinamento de modelo vai levar muito tempo. Pressione Ctrl+C e cancele o treinamento local em vez de esperar a conclusão do job.

Se você quiser fazer mais testes, também é possível executar diretamente a imagem criada acima, sem reempacotamento.

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 personalizada do Docker com base no seu código.
  • Envie a imagem para o Container Registry
  • Inicia um CustomJob com base na imagem.

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

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

Especifique a imagem da GPU do TensorFlow pré-criada do Vertex AI Training como a imagem de base do 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, o comando vai criar uma imagem Docker personalizada com base no código de treinamento. A imagem base é o contêiner pré-criado do Vertex AI Training definido como BASE_GPU_IMAGE. O recurso de empacotamento automático instalará com pip 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. Define a configuração do pool de workers usada pelo job personalizado. Especifique 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.

Veja alguns dos principais campos dessa especificação:

  • machine-type (obrigatório): o tipo de máquina. Clique aqui para ver os tipos compatíveis.
  • replica-count: o número de réplicas de workers a serem usadas para este pool de workers. Por padrão, o valor é 1.
  • accelerator-type: o tipo de GPUs. Clique aqui para ver os tipos compatíveis. Neste exemplo, especificamos uma GPU NVIDIA Tesla V100.
  • accelerator-count: o número de GPUs a serem usadas por cada VM no pool de workers. Por padrão, o valor é 1.
  • executor-image-uri: o URI de uma imagem de contêiner que executará o pacote fornecido. Ele é 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 na execução do comando local, você verá a imagem Docker sendo criada e o job de treinamento será iniciado. Exceto em vez de ver a saída do job de treinamento, você verá a seguinte mensagem confirmando que o job de treinamento foi iniciado. Observe que, na primeira vez que você executar o comando custom-jobs create, pode levar alguns minutos para a imagem ser criada e enviada.

training_started

Volte à seção "Treinamento da Vertex AI" do Console do Cloud e, em JOBS PERSONALIZADOS, seu job será executado.

training_job

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

Depois da conclusão, você verá os seguintes artefatos de modelo salvos no diretório model_output do bucket.

model_output

Parabéns! 🎉

Você aprendeu a usar a Vertex AI para:

  • Conteinerizar e executar o código de treinamento 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