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 personalizados 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 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. 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. 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. Nesse caso, temos um problema de classificação binária (positivo ou negativo), então especificamos num_labels=2 para esse classificador.

6. Conteinerizar e executar o código de treinamento localmente

Use 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 execute 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 treinamento da Vertex AI e pode ajudar a depurar problemas com o código antes de realizar o treinamento personalizado na Vertex AI.

No job de treinamento, exportamos 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 executamos o job de treinamento personalizado no treinamento da Vertex AI, usamos 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 seguinte comando para definir o URI de uma imagem do 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 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 base, então precisamos fornecê-las como requisitos. Para 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 nossa 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

A imagem do Docker vai ser 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 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 do 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 reembalar.

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:

  • 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 acesse um nível acima do diretório autopkg-codelab.

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

Especifique a imagem pré-criada do TensorFlow para GPU do treinamento da Vertex AI 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, esse comando vai criar uma imagem personalizada do Docker com base no código de treinamento. A imagem de base é o contêiner pré-criado da Vertex AI Training que definimos como BASE_GPU_IMAGE. O recurso de autoempacotamento 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. Isso define a configuração do pool de workers usada pelo job personalizado. Você pode 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 único pool de workers, porque o código de treinamento não está configurado para treinamento distribuído.

Estes são alguns dos principais campos desta 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 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 que cada VM no pool de workers vai usar. Por padrão, o valor é 1.
  • executor-image-uri: o URI de uma imagem de contêiner que vai executar o pacote fornecido. Ela está definida 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 que será executado no pacote fornecido.

Assim como quando você executou o comando local, a imagem do Docker será criada e o job de treinamento será iniciado. No entanto, em vez de ver a saída do job de treinamento, você vai ver 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 "Treinamento da Vertex AI" do console do Cloud e, em JOBS PERSONALIZADOS, você vai encontrar o job em execução.

training_job

A conclusão do job leva cerca de 20 minutos.

Depois de concluído, os seguintes artefatos de modelo salvos vão aparecer 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
  • Envie jobs de treinamento para o treinamento da Vertex AI com o 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