Vertex AI: treinamento de vários workers e aprendizado por transferência com o TensorFlow

1. Visão geral

Neste laboratório, você usará a Vertex AI para executar um job de treinamento de vários workers para um modelo do TensorFlow.

Conteúdo do laboratório

Você aprenderá como realizar as seguintes tarefas:

• Modificar o código do aplicativo de treinamento para treinamento de vários workers
• Configurar e iniciar um job de treinamento de vários workers na IU da Vertex AI
• Configurar e iniciar um job de treinamento de vários workers com o SDK do Vertex

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

2. Introdução ao 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 os dois tipos em uma única API, junto com outros produtos novos. 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 inclui vários produtos diferentes para dar suporte a fluxos de trabalho integrais de ML. Este laboratório aborda os produtos destacados abaixo: Treinamento e Workbench

3. Visão geral do caso de uso

Neste laboratório, você usará o aprendizado por transferência para treinar um modelo de classificação de imagens no conjunto de dados do Mansava em Conjuntos de dados do TensorFlow. A arquitetura que você usará é um modelo do ResNet50 da biblioteca tf.keras.applications pré-treinado no conjunto de dados do Imagenet.

Por que usar o treinamento distribuído?

Se você tiver apenas uma GPU, o TensorFlow usará esse acelerador para acelerar o treinamento do modelo sem nenhum trabalho extra da sua parte. No entanto, se você quiser ter um impulso adicional no uso de várias GPUs em uma ou várias máquinas (cada uma com várias GPUs), será necessário usar tf.distribute, que é a biblioteca do TensorFlow para executar uma computação em vários dispositivos. Um dispositivo se refere a uma CPU ou acelerador, como GPUs ou TPUs, em alguma máquina em que o TensorFlow possa executar operações.

A maneira mais simples de começar com o treinamento distribuído é uma máquina com vários dispositivos de GPU. Uma estratégia de distribuição do TensorFlow do módulo tf.distribute gerenciará a coordenação da distribuição de dados e das atualizações de gradiente em todas as GPUs. Se você já domina o treinamento de um único host e quer escalonar ainda mais, adicionar várias máquinas ao cluster pode ajudar a melhorar ainda mais o desempenho. É possível usar um cluster de máquinas que são apenas de CPU ou que têm uma ou mais GPUs. Este laboratório aborda o último caso e demonstra como usar o MultiWorkerMirroredStrategy para distribuir o treinamento de um modelo do TensorFlow em várias máquinas na Vertex AI.

MultiWorkerMirroredStrategy é uma estratégia de paralelismo síncrono que pode ser usada com apenas algumas alterações no código. Uma cópia do modelo é criada em cada dispositivo no cluster. As próximas atualizações de gradientes ocorrerão de maneira síncrona. Isso significa que cada dispositivo de worker calcula a passagem do modelo para frente e para trás em uma fração diferente dos dados de entrada. Os gradientes calculados de cada uma dessas frações são agregados em todos os dispositivos de uma máquina e em todas as máquinas do cluster e reduzidos (geralmente uma média) em um processo conhecido como "redução total". Depois, o otimizador executa as atualizações de parâmetro com esses gradientes reduzidos, mantendo os dispositivos em sincronia. Para saber mais sobre o treinamento distribuído com o TensorFlow, confira o vídeo abaixo:

4. configure o ambiente

Para executar este codelab, você 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. Você precisará dele para criar sua instância de notebook.

Etapa 2: ativar a API Container Registry

Navegue até o Container Registry e selecione Ativar. Use-o para criar um contêiner de job de treinamento personalizado.

Etapa 3: ativar a API Vertex AI

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

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

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

Ative a API Notebooks se ela ainda não tiver sido ativada.

Após a ativação, clique em NOTEBOOK GERENCIADO:

Em seguida, selecione NOVO NOTEBOOK.

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

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.

Em Segurança, selecione "Ativar terminal" se essa opção ainda não estiver ativada.

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

Em seguida, clique em Criar. O provisionamento da instância levará alguns minutos.

Após a criação da instância, selecione Abrir o JupyterLab.

Na primeira vez que usar uma nova instância, você receberá uma solicitação para autenticar. Siga as etapas na IU para isso.

5. Conteinerizar o código do aplicativo de treinamento

Para enviar esse job de treinamento à Vertex, coloque o código do aplicativo de treinamento em um contêiner do Docker e envie esse contêiner ao Google Container Registry. Com essa abordagem, é possível treinar um modelo criado com qualquer framework.

Para começar, no menu "Launcher", abra uma janela de terminal na instância do notebook:

Crie um diretório chamado cassava e use cd nele:

mkdir cassava
cd cassava

Etapa 1: criar um Dockerfile

A primeira etapa para inserir o código em contêiner é criar um Dockerfile. No Dockerfile, inclua todos os comandos necessários para executar a imagem. Ele instalará todas as bibliotecas necessárias e configurará o ponto de entrada do código de treinamento.

No seu terminal, crie um Dockerfile vazio:

touch Dockerfile

Abra o Dockerfile e copie o seguinte nele:

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"]

Esse Dockerfile usa a imagem do Docker do GPU Enterprise Contêiner do TensorFlow 2.7. Os contêineres de aprendizado profundo no Google Cloud vêm com muitos frameworks comuns de ML e ciência de dados pré-instalados. Após o download, o Dockerfile configura o ponto de entrada para o código de treinamento. Você ainda não criou esses arquivos. Na próxima etapa, você adicionará o código para treinar e ajustar o modelo.

Etapa 2: criar um bucket do Cloud Storage

Neste job de treinamento, você vai exportar o modelo treinado do TensorFlow para um bucket do Cloud Storage. No terminal, execute o seguinte para definir uma variável de ambiente para seu projeto, substituindo your-cloud-project pelo ID do projeto:

Execute o gcloud config list --format 'value(core.project)' no seu terminal para conseguir o ID do projeto.

PROJECT_ID='your-cloud-project'

Em seguida, execute o comando a seguir no Terminal para criar um novo bucket no projeto:

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

Etapa 3: adicionar o código de treinamento de modelo

No seu terminal, execute o seguinte comando para criar um diretório para o código de treinamento e um arquivo Python em que você adicionará o código:

mkdir trainer
touch trainer/task.py

Agora você tem o seguinte no seu diretório cassava/:

+ Dockerfile
+ trainer/
    + task.py

Em seguida, abra o arquivo task.py que você acabou de criar e copie o código abaixo. Será necessário substituir {your-gcs-bucket} pelo nome do bucket do Cloud Storage que você acabou de criar.

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 criar o contêiner, vamos analisar o código mais detalhadamente, que usa MultiWorkerMirroredStrategy da API tf.distribute.Strategy.

O MultiWorkerMirroredStrategy foi projetado para distribuir o treinamento entre várias máquinas em um cluster, cada uma com possíveis GPUs. Se você quiser distribuir o treinamento em várias GPUs apenas em um dispositivo, consulte MirroredStrategy.

Há alguns componentes no código que são necessários para que seu código funcione com MultiWorkerMirroredStrategy.

1. Os dados precisam ser fragmentados, o que significa que cada worker recebe um subconjunto de todo o conjunto de dados. Portanto, em cada etapa, um tamanho global de lote de elementos de conjuntos de dados não sobrepostos será processado por cada worker. Essa fragmentação acontece automaticamente com tf.data.experimental.AutoShardPolicy, que pode ser definido como FILE ou DATA. Neste exemplo, a função create_dataset() define AutoShardPolicy como DATA porque o conjunto de dados da mandioca não foi transferido como vários arquivos. No entanto, se você não tiver definido a política como DATA, a política AUTO padrão será iniciada e o resultado final será o mesmo. Saiba mais sobre a fragmentação de conjuntos de dados com o MultiWorkerMirroredStrategy.
2. Na função main(), o objeto MultiWorkerMirroredStrategy é criado. Depois, envolva a criação das variáveis de modelo no escopo da estratégia. Essa etapa crucial informa ao TensorFlow quais variáveis precisam ser espelhadas nas réplicas.
3. O tamanho do lote é escalonado verticalmente pelo num_replicas_in_sync. Isso garante que cada réplica processe o mesmo número de exemplos em cada etapa. O escalonamento do tamanho do lote é uma prática recomendada ao usar estratégias de paralelismo síncrono de dados no TensorFlow.
4. Salvar seu modelo é um pouco mais complicado no caso de vários workers porque o destino precisa ser diferente para cada um deles. O worker principal será salvo no diretório do modelo desejado, e os outros workers vão salvar o modelo em diretórios temporários. É importante que esses diretórios temporários sejam exclusivos para evitar que vários workers gravem no mesmo local. Salvar pode conter operações coletivas, o que significa que todos os workers devem salvar, e não somente o chefe. As funções _is_chief(), _get_temp_dir(), write_filepath() e main() incluem um código boilerplate que ajuda a salvar o modelo.

Se você usou MultiWorkerMirroredStrategy em um ambiente diferente, talvez tenha configurado a variável de ambiente TF_CONFIG. A Vertex AI define TF_CONFIG automaticamente para você, então não é necessário definir essa variável em cada máquina do cluster.

Etapa 4: criar o contêiner

No terminal, execute o seguinte para definir uma variável de ambiente para seu projeto, substituindo your-cloud-project pelo ID do projeto:

Execute o gcloud config list --format 'value(core.project)' no seu terminal para conseguir o ID do projeto.

PROJECT_ID='your-cloud-project'

Defina uma variável com o URI da imagem do contêiner no Google Container Registry:

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

Em seguida, crie o contêiner executando a seguinte linha a partir da raiz do diretório cassava:

docker build ./ -t $IMAGE_URI

Por fim, envie-o para o Google Container Registry:

docker push $IMAGE_URI

Com o contêiner enviado para o Container Registry, é possível iniciar um job de treinamento.

6. Executar um job de treinamento de vários workers na Vertex AI

Este laboratório usa treinamento personalizado com um contêiner personalizado no Google Container Registry, mas você também pode executar um job de treinamento com os contêineres pré-criados.

Para começar, navegue até a seção Training na seção "Vertex" do Console do Cloud:

Etapa 1: configurar o job de treinamento

Clique em Criar para inserir os parâmetros do job de treinamento.

• Em Conjunto de dados, selecione Nenhum conjunto de dados gerenciado.
• Em seguida, selecione Treinamento personalizado (avançado) como o método de treinamento e clique em Continuar.
• Insira multiworker-cassava (ou o nome que você quer atribuir ao modelo) em Nome do modelo
• Clique em Continuar

Na etapa "Configurações do contêiner", selecione Contêiner personalizado:

Na primeira caixa (Imagem do contêiner), insira o valor da variável IMAGE_URI da seção anterior. Precisa ser: gcr.io/your-cloud-project/multiworker:cassava, com o ID do seu projeto. Deixe os outros campos em branco e clique em Continuar.

Pule a etapa dos hiperparâmetros. Para isso, clique em Continuar novamente.

Etapa 2: configurar o cluster de computação

A Vertex AI oferece quatro pools de workers para cobrir os diferentes tipos de tarefas de máquina.

O pool de workers 0 configura o principal, o chefe, o programador ou o "mestre". Em MultiWorkerMirroredStrategy, todas as máquinas são designadas como workers, que são as máquinas físicas em que o cálculo replicado é executado. Além de cada máquina ser um worker, é necessário que haja um worker que assuma um trabalho extra, como salvar checkpoints e gravar arquivos de resumo no TensorBoard. Esta máquina é conhecida como o chefe. Só há um chefe de workers para que a contagem do pool de workers 0 seja sempre 1.

Em Computação e preços, mantenha a região selecionada e configure o Pool de workers 0 da seguinte maneira:

É possível adicionar GPUs, mas lembre-se de que, para MultiWorkerMirroredStrategy, cada máquina do cluster precisa ter o mesmo número de GPUs. A adição de GPUs também aumentará o custo do job de treinamento.

O pool de workers 1 é onde você configura os workers no cluster.

Configure o Pool de workers 1 da seguinte maneira:

Agora o cluster está configurado para ter duas máquinas somente de CPU. Quando o código do aplicativo de treinamento é executado, o MultiWorkerMirroredStrategy distribui o treinamento nas duas máquinas.

MultiWorkerMirroredStrategy tem apenas os tipos de tarefa do chefe e do worker. Portanto, não é necessário configurar outros pools de workers. No entanto, se você usar a ParameterServerStrategy do TensorFlow, precisará configurar os servidores de parâmetros no Pool de workers 2. Se você quiser adicionar um avaliador ao cluster, configure essa máquina no Pool de workers 3.

Clique em Iniciar treinamento para iniciar o job de ajuste de hiperparâmetros. Na seção "Training" do console, na guia TRAINING PIPELINES, você verá o job recém-lançado:

Esse job levará de 30 a 35 minutos para ser concluído.

Parabéns! 🎉

Você aprendeu a usar a Vertex AI para:

• iniciar um job de treinamento de vários workers para o código de treinamento fornecido em um contêiner personalizado. Neste exemplo, você usou um modelo do TensorFlow, mas é possível treinar um modelo criado com qualquer framework usando contêineres personalizados ou integrados.

Para saber mais sobre diferentes partes do Vertex, confira a documentação.

7. [Opcional] Usar o SDK do Vertex

A seção anterior mostrou como iniciar o job de treinamento usando a IU. Nesta seção, você verá uma maneira alternativa de enviar o job de treinamento usando a API Vertex Python.

Volte para a instância do notebook e crie um notebook do TensorFlow 2 no acesso rápido:

Importe o SDK da Vertex AI.

from google.cloud import aiplatform

Para iniciar o job de treinamento de vários workers, primeiro defina a especificação do pool de workers. O uso de GPUs na especificação é completamente opcional e é possível remover accelerator_type e accelerator_count se quiser um cluster somente de CPU, conforme mostrado na seção 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"}
      }
]

Em seguida, crie e execute um CustomJob. Será necessário substituir {YOUR_BUCKET} por um bucket no seu projeto para preparo. É possível usar o mesmo bucket criado 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()

Na seção "Treinamento" do console, na guia JOBS PERSONALIZADOS, você verá seu job de treinamento:

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 "Stop" na seção "Vertex AI Workbench" do console. Se quiser excluir o notebook completamente, clique no botão 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":