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

1. Visão geral

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

Conteúdo do laboratório

Você vai aprender a:

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

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

2. 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 inclui vários produtos diferentes para dar suporte a fluxos de trabalho integrais de ML. Os produtos destacados abaixo são o foco deste laboratório: Treinamentos e Workbench.

3. Visão geral do caso de uso

Neste laboratório, você vai usar o aprendizado por transferência para treinar um modelo de classificação de imagens no conjunto de dados de mandioca dos conjuntos de dados do TensorFlow. Você usará um modelo ResNet50 da biblioteca tf.keras.applications pré-treinado no conjunto de dados Imagenet.

Por que usar treinamento distribuído?

Se você tiver apenas uma GPU, o TensorFlow vai usar esse acelerador no treinamento do modelo sem nenhum trabalho extra da sua parte. No entanto, se você quiser aumentar ainda mais o uso de várias GPUs em uma ou várias máquinas (cada uma podendo ter várias GPUs), use a tf.distribute, que é a biblioteca do TensorFlow para executar 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 pode executar operações.

A maneira mais simples de começar o treinamento distribuído é uma única máquina com vários dispositivos GPU. Uma estratégia de distribuição do TensorFlow do módulo tf.distribute gerenciará a coordenação da distribuição de dados e atualizações de gradiente em todas as GPUs. Se você domina o treinamento de host único e quer escalonar ainda mais, adicionar várias máquinas ao seu cluster pode ajudar a otimizar ainda mais o desempenho. É possível usar um cluster de máquinas que sejam apenas CPU ou que tenham 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.

A MultiWorkerMirroredStrategy é uma estratégia síncrona de carregamento em paralelo de dados 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 as passagens anteriores e posteriores do modelo 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 em uma máquina e em todas as máquinas no cluster e reduzidos (geralmente, uma média) em um processo conhecido como redução total. Em seguida, o otimizador realiza as atualizações dos parâmetros com esses gradientes reduzidos, mantendo os dispositivos sincronizados. Para saber mais sobre treinamento distribuído com o TensorFlow, confira o vídeo abaixo:

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. Você vai precisar disso para criar sua instância de notebook.

Etapa 2: 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 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 vai levar alguns minutos.

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

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

5. fazer a conteinerização do código do aplicativo de treinamento

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

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

Crie um novo diretório chamado cassava e coloque cd nele:

mkdir cassava
cd cassava

Etapa 1: criar um Dockerfile

A primeira etapa na conteinerização de seu código é a criação de um Dockerfile. Nesse Dockerfile, você vai incluir todos os comandos necessários à execução da imagem. Ele vai 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"]

Este Dockerfile usa a imagem do Docker do TensorFlow Enterprise 2.7 GPU no Deep Learning Container. O componente Deep Learning Containers no Google Cloud vem com vários frameworks de ciência de dados já instalados. Após fazer download dessa imagem, este Dockerfile configura o ponto de entrada para o código de treinamento. Você ainda não criou esses arquivos – na próxima etapa, você vai 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 do TensorFlow treinado para o 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.

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

PROJECT_ID='your-cloud-project'

Depois 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 do modelo

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 cassava/:

+ Dockerfile
+ trainer/
    + task.py

Depois abra o arquivo task.py que você acabou de criar e copie o código abaixo. Substitua {your-gcs-bucket} pelo nome do bucket do Cloud Storage recém-criado.

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, que usa o MultiWorkerMirroredStrategy da API tf.distribute.Strategy.

O MultiWorkerMirroredStrategy foi projetado para distribuir treinamento em várias máquinas de um cluster, cada uma com várias GPUs. Se você quiser distribuir o treinamento para várias GPUs em um único 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 de lote global de elementos do conjunto 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 não é feito o download do conjunto de dados de mandioca como vários arquivos. No entanto, se você não a definir como DATA, a política AUTO padrão será ativada e o resultado será o mesmo. Saiba mais sobre a fragmentação de conjuntos de dados com o MultiWorkerMirroredStrategy aqui.
2. Na função main(), o objeto MultiWorkerMirroredStrategy é criado. Depois, envolva a criação das variáveis de modelo no escopo da estratégia. Esta 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. É um pouco mais complicado salvar seu modelo no caso de vários workers porque o destino precisa ser diferente para cada um. O worker CHIEF vai ser salvo no diretório do modelo desejado, enquanto os demais vão salvar o modelo em diretórios temporários. É importante que esses diretórios temporários sejam únicos para impedir que vários workers gravem no mesmo local. O salvamento pode conter operações coletivas: todos os workers precisam salvar e não apenas o CHIEF. As funções _is_chief(), _get_temp_dir(), write_filepath() e main() incluem 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 a variável TF_CONFIG automaticamente para que você não precise fazer isso em cada máquina no cluster.

Etapa 4: criar o contêiner

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.

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 seu contêiner no Google Container Registry:

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

Configure o Docker.

gcloud auth configure-docker

Agora execute o comando a seguir na raiz do diretório cassava para criar o diretório:

docker build ./ -t $IMAGE_URI

Por fim, envie para o Google Container Registry:

docker push $IMAGE_URI

Depois que o contêiner for enviado para o Container Registry, estará tudo pronto para você iniciar o job de treinamento.

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

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

Para começar, navegue até a seção Treinamento na seção da Vertex do console do Cloud:

Etapa 1: configurar job de treinamento

Clique em Criar e insira os parâmetros para seu job de treinamento.

• Em Conjunto de dados, selecione Sem conjunto de dados gerenciado.
• Selecione Treinamento personalizado (avançado) como seu método de treinamento e clique em Continuar.
• Insira multiworker-cassava (ou como quiser nomear seu modelo) em Nome do modelo.
• Clique em Continuar.

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

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

Pule a etapa de hiperparâmetros clicando em Continuar novamente.

Etapa 2: configurar cluster de computação

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

O pool de workers 0 configura o principal, CHIEF, programador ou "mestre". No 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, é preciso que um deles fique encarregado do trabalho extra, como salvar checkpoints e gravar arquivos de resumo no TensorBoard. Essa máquina é conhecida como CHIEF. Como há apenas um worker CHIEF, a contagem de workers no pool 0 sempre vai ser 1.

Em Computação e preços, deixe a região selecionada como está e configure o Pool de workers 0 da seguinte forma:

Você pode optar por adicionar GPUs, mas lembre-se de que, para MultiWorkerMirroredStrategy, cada máquina no cluster deve ter o mesmo número de GPUs. Adicionar GPUs também aumentará o custo do job de treinamento.

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

Configure o pool de workers 1 como flores:

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

MultiWorkerMirroredStrategy tem apenas os tipos de tarefa CHIEF e worker. Por isso, não é necessário configurar mais pools de workers. No entanto, se você usasse o ParameterServerStrategy do TensorFlow, configuraria os servidores de parâmetros no pool de workers 2. E, para adicionar um avaliador ao cluster, basta configurar essa máquina no pool de workers 3.

Clique em Iniciar treinamento para começar o job de ajuste de hiperparâmetros. Na seção "Treinamento" do console, na guia PIPELINES DE TREINAMENTO, você vai encontrar o job recém-iniciado.

O job de treinamento leva cerca de 30 a 35 minutos para ser concluído.

Parabéns! 🎉

Você aprendeu a usar a Vertex AI para:

• Iniciar um job de treinamento com 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 estrutura usando contêineres personalizados ou integrados.

Para saber mais sobre as diferentes partes da Vertex, consulte a documentação.

7. use o SDK Vertex [Opcional]

A seção anterior mostrou como iniciar o job de treinamento pela interface. Nesta seção, você vai encontrar uma maneira alternativa de enviar o job de treinamento usando a API Vertex Python.

Retorne à sua instância do notebook e crie um notebook do TensorFlow 2 na tela de início:

Importe o SDK da Vertex AI.

from google.cloud import aiplatform

Para iniciar o job de treinamento com vários workers, primeiro é preciso definir a especificação do pool de workers. O uso de GPUs na especificação é totalmente opcional, e você pode remover accelerator_type e accelerator_count se quiser um cluster somente para CPU, como 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. Você vai precisar substituir {YOUR_BUCKET} por um bucket no 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, seu job de treinamento vai aparecer na guia JOBS PERSONALIZADOS:

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

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