Treinamento e ajuste de hiperparâmetros em um modelo PyTorch na AI Platform do Cloud

1. Visão geral

Neste laboratório, você vai conhecer um fluxo de trabalho completo de treinamento de ML no Google Cloud usando o PyTorch para criar seu modelo. Em um ambiente do Notebooks da IA do Google Cloud, você vai aprender a empacotar seu job de treinamento para executá-lo no AI Platform Training com ajuste de hiperparâmetros.

Conteúdo do laboratório

Você vai aprender a:

• Criar uma instância do AI Platform Notebooks
• Criar um modelo do PyTorch
• Treinar o modelo com ajuste de hiperparâmetros no AI Platform Training

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

2. 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 Cloud AI Platform Models

Acesse a seção "Modelos do AI Platform" do console do Cloud e clique em "Ativar", se ainda não tiver feito isso.

Etapa 2: 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 3: criar uma instância do AI Platform Notebooks

Acesse a seção AI Platform Notebooks do console do Cloud e clique em Nova instância. Em seguida, selecione o tipo de instância PyTorch mais recente (sem GPUs):

Use as opções padrão ou dê um nome personalizado, se quiser, e clique em Criar. Quando a instância tiver sido criada, selecione Abrir o JupyterLab:

Em seguida, abra uma instância de notebook Python 3 na tela de início:

Tudo pronto para começar!

Etapa 5: importar pacotes Python

No restante deste codelab, execute todos os snippets de código do seu notebook do Jupyter.

Na primeira célula do notebook, adicione as importações a seguir e execute a célula. Para executar, pressione o botão de seta para a direita no menu superior ou pressione command-enter:

import datetime
import numpy as np
import os
import pandas as pd
import time

Você vai notar que não estamos importando o PyTorch aqui. Isso acontece porque estamos executando o job de treinamento no AI Platform Training, não na instância do notebook.

3. Criar um pacote para o job de treinamento

Para executar nosso job de treinamento no AI Platform Training, precisamos do código de treinamento empacotado localmente na instância do Notebooks e de um bucket do Cloud Storage para armazenar recursos do job. Primeiro, vamos criar um bucket de armazenamento. Pule esta etapa se você já tiver um.

Etapa 1: criar um bucket do Cloud Storage para nosso modelo

Primeiro, vamos definir algumas variáveis de ambiente que vamos usar no restante do codelab. Preencha os valores abaixo com o nome do seu projeto na nuvem do Google Cloud e o nome do bucket do Cloud Storage que você quer criar (precisa ser globalmente exclusivo):

# Update these to your own GCP project, model, and version names
GCP_PROJECT = 'your-gcp-project'
BOCKET_URL = 'gs://storage_bucket_name'

Agora está tudo pronto para criar um bucket de armazenamento, que vamos apontar quando iniciarmos o job de treinamento.

Execute este comando gsutil no notebook para criar um bucket:

!gsutil mb $BUCKET_URL

Etapa 2: criar os arquivos iniciais para nosso pacote Python

Para executar um job de treinamento no AI Platform, precisamos configurar nosso código como um pacote do Python. Ele consiste em um arquivo setup.py no diretório raiz que especifica as dependências de pacotes externos, um subdiretório com o nome do pacote (aqui vamos chamar de trainer/) e um arquivo __init__.py vazio dentro desse subdiretório.

Primeiro, vamos escrever o arquivo setup.py. Estamos usando os comandos mágicos %%writefile do iPython para salvar o arquivo na nossa instância. Aqui, especificamos três bibliotecas externas que vamos usar no nosso código de treinamento: PyTorch, Scikit-learn e Pandas:

%%writefile setup.py
from setuptools import find_packages
from setuptools import setup

REQUIRED_PACKAGES = ['torch>=1.5', 'scikit-learn>=0.20', 'pandas>=1.0']

setup(
    name='trainer',
    version='0.1',
    install_requires=REQUIRED_PACKAGES,
    packages=find_packages(),
    include_package_data=True,
    description='My training application package.'
)

Em seguida, vamos criar o diretório trainer/ e o arquivo init.py vazio dentro dele. O Python usa esse arquivo para reconhecer que este é um pacote:

!mkdir trainer
!touch trainer/__init__.py

Agora está tudo pronto para começar a criar o job de treinamento.

4. Visualizar o conjunto de dados

O foco deste laboratório é nas ferramentas para treinar modelos, mas vamos dar uma olhada rápida no conjunto de dados que usaremos para treinar nosso modelo. Vamos usar o conjunto de dados de natalidade disponível no BigQuery. Ele contém dados de nascimento dos EUA de várias décadas. Vamos usar algumas colunas do conjunto de dados para prever o peso de um bebê ao nascer. O conjunto de dados original é muito grande, e vamos usar um subconjunto dele que disponibilizamos para você em um bucket do Cloud Storage.

Etapa 1: fazer o download do conjunto de dados de natalidade do BigQuery

Vamos baixar a versão do conjunto de dados que disponibilizamos para você no Cloud Storage em um DataFrame do Pandas e visualizar.

natality = pd.read_csv('https://storage.googleapis.com/ml-design-patterns/natality.csv')
natality.head()

Esse conjunto de dados tem pouco menos de 100.000 linhas. Vamos usar cinco recursos para prever o peso de nascimento de um bebê: idade da mãe e do pai, semanas de gestação, ganho de peso da mãe em libras e sexo do bebê representado como um booleano.

5. Definir o job de treinamento com ajuste de hiperparâmetros

Vamos escrever nosso script de treinamento em um arquivo chamado model.py no subdiretório trainer/ que criamos anteriormente. Nosso job de treinamento será executado no AI Platform Training e também usará o serviço de ajuste de hiperparâmetros do AI Platform para encontrar os hiperparâmetros ideais para nosso modelo usando a otimização bayesiana.

Etapa 1: criar o script de treinamento

Primeiro, vamos criar o arquivo Python com nosso script de treinamento. Depois, vamos analisar o que está acontecendo nele. Executar esse comando %%writefile vai gravar o código do modelo em um arquivo Python local:

%%writefile trainer/model.py
import argparse
import hypertune
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim

from sklearn.utils import shuffle
from sklearn.preprocessing import LabelEncoder
from sklearn.preprocessing import normalize

def get_args():
    """Argument parser.
    Returns:
        Dictionary of arguments.
    """
    parser = argparse.ArgumentParser(description='PyTorch MNIST')
    parser.add_argument('--job-dir',  # handled automatically by AI Platform
                        help='GCS location to write checkpoints and export ' \
                             'models')
    parser.add_argument('--lr',  # Specified in the config file
                        type=float,
                        default=0.01,
                        help='learning rate (default: 0.01)')
    parser.add_argument('--momentum',  # Specified in the config file
                        type=float,
                        default=0.5,
                        help='SGD momentum (default: 0.5)')
    parser.add_argument('--hidden-layer-size',  # Specified in the config file
                        type=int,
                        default=8,
                        help='hidden layer size')
    args = parser.parse_args()
    return args

def train_model(args):
    # Get the data
    natality = pd.read_csv('https://storage.googleapis.com/ml-design-patterns/natality.csv')
    natality = natality.dropna()
    natality = shuffle(natality, random_state = 2)
    natality.head()

    natality_labels = natality['weight_pounds']
    natality = natality.drop(columns=['weight_pounds'])


    train_size = int(len(natality) * 0.8)
    traindata_natality = natality[:train_size]
    trainlabels_natality = natality_labels[:train_size]

    testdata_natality = natality[train_size:]
    testlabels_natality = natality_labels[train_size:]

    # Normalize and convert to PT tensors
    normalized_train = normalize(np.array(traindata_natality.values), axis=0)
    normalized_test = normalize(np.array(testdata_natality.values), axis=0)

    train_x = torch.Tensor(normalized_train)
    train_y = torch.Tensor(np.array(trainlabels_natality))

    test_x = torch.Tensor(normalized_test)
    test_y = torch.Tensor(np.array(testlabels_natality))

    # Define our data loaders
    train_dataset = torch.utils.data.TensorDataset(train_x, train_y)
    train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True)

    test_dataset = torch.utils.data.TensorDataset(test_x, test_y)
    test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=128, shuffle=False)

    # Define the model, while tuning the size of our hidden layer
    model = nn.Sequential(nn.Linear(len(train_x[0]), args.hidden_layer_size),
                          nn.ReLU(),
                          nn.Linear(args.hidden_layer_size, 1))
    criterion = nn.MSELoss()

    # Tune hyperparameters in our optimizer
    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
    epochs = 20
    for e in range(epochs):
        for batch_id, (data, label) in enumerate(train_dataloader):
            optimizer.zero_grad()
            y_pred = model(data)
            label = label.view(-1,1)
            loss = criterion(y_pred, label)
            
            loss.backward()
            optimizer.step()


    val_mse = 0
    num_batches = 0
    # Evaluate accuracy on our test set
    with torch.no_grad():
        for i, (data, label) in enumerate(test_dataloader):
            num_batches += 1
            y_pred = model(data)
            mse = criterion(y_pred, label.view(-1,1))
            val_mse += mse.item()


    avg_val_mse = (val_mse / num_batches)

    # Report the metric we're optimizing for to AI Platform's HyperTune service
    # In this example, we're mimizing error on our test set
    hpt = hypertune.HyperTune()
    hpt.report_hyperparameter_tuning_metric(
        hyperparameter_metric_tag='val_mse',
        metric_value=avg_val_mse,
        global_step=epochs        
    )

def main():
    args = get_args()
    print('in main', args)
    train_model(args)

if __name__ == '__main__':
    main()

O job de treinamento consiste em duas funções em que a maior parte do trabalho está acontecendo.

• get_args(): analisa os argumentos de linha de comando que vamos transmitir ao criar o job de treinamento, além dos hiperparâmetros que queremos que o AI Platform otimize. Neste exemplo, nossa lista de argumentos inclui apenas os hiperparâmetros que vamos otimizar: a taxa de aprendizado, o momentum e o número de neurônios na camada escondida do modelo.
• train_model(): aqui baixamos os dados para um DataFrame do Pandas, normalizamos e convertemos em tensores do PyTorch e definimos nosso modelo. Para criar nosso modelo, usamos a API nn.Sequential do PyTorch, que permite definir o modelo como uma pilha de camadas:

model = nn.Sequential(nn.Linear(len(train_x[0]), args.hidden_layer_size),
                      nn.ReLU(),
                      nn.Linear(args.hidden_layer_size, 1))

Observe que, em vez de fixar no código o tamanho da camada escondida do nosso modelo, estamos transformando isso em um hiperparâmetro que o AI Platform vai ajustar para nós. Saiba mais sobre isso na próxima seção.

Etapa 2: usar o serviço de ajuste de hiperparâmetros da AI Platform

Em vez de tentar manualmente diferentes valores de hiperparâmetros e treinar novamente nosso modelo a cada vez, vamos usar o serviço de otimização de hiperparâmetros do AI Platform do Google Cloud. Se configurarmos nosso job de treinamento com argumentos de hiperparâmetros, o AI Platform usará a otimização bayesiana para encontrar os valores ideais para os hiperparâmetros especificados.

No ajuste de hiperparâmetros, um único teste consiste em uma execução de treinamento do nosso modelo com uma combinação específica de valores de hiperparâmetros. Dependendo de quantos testes executamos, o AI Platform usa os resultados dos testes concluídos para otimizar os hiperparâmetros selecionados para os futuros. Para configurar o ajuste de hiperparâmetros, precisamos transmitir um arquivo de configuração ao iniciar o job de treinamento com alguns dados sobre cada um dos hiperparâmetros que estamos otimizando.

Em seguida, crie o arquivo de configuração localmente:

%%writefile config.yaml
trainingInput:
  hyperparameters:
    goal: MINIMIZE
    maxTrials: 10
    maxParallelTrials: 5
    hyperparameterMetricTag: val_mse
    enableTrialEarlyStopping: TRUE
    params:
    - parameterName: lr
      type: DOUBLE
      minValue: 0.0001
      maxValue: 0.1
      scaleType: UNIT_LINEAR_SCALE
    - parameterName: momentum
      type: DOUBLE
      minValue: 0.0
      maxValue: 1.0
      scaleType: UNIT_LINEAR_SCALE
    - parameterName: hidden-layer-size
      type: INTEGER
      minValue: 8
      maxValue: 32
      scaleType: UNIT_LINEAR_SCALE

Para cada hiperparâmetro, especificamos o tipo, o intervalo de valores que queremos pesquisar e a escala em que aumentar o valor em diferentes testes.

No início do job, também especificamos a métrica que estamos otimizando. Observe que, no final da função train_model() acima, informamos essa métrica ao AI Platform sempre que um teste é concluído. Aqui, estamos minimizando o erro médio quadrático do nosso modelo. Por isso, queremos usar os hiperparâmetros que resultam no menor erro médio quadrático para nosso modelo. O nome dessa métrica (val_mse) corresponde ao nome usado para informar quando chamamos report_hyperparameter_tuning_metric() ao final de um teste.

Observação: embora não estejamos executando nosso job de treinamento de modelo diretamente na instância do notebook, você pode copiar o código de treinamento do modelo para uma nova célula e testar no notebook. Basta substituir todas as referências a args por valores de hiperparâmetros fixados no código. Como você já criou uma instância do PyTorch, está tudo pronto.

6. Executar um job de treinamento no AI Platform

Nesta seção, vamos iniciar nosso job de treinamento de modelo com ajuste de hiperparâmetros na AI Platform.

Etapa 1: definir algumas variáveis de ambiente

Primeiro, vamos definir algumas variáveis de ambiente que usaremos para iniciar o job de treinamento. Se você quiser executar o job em uma região diferente, atualize a variável REGION abaixo:

MAIN_TRAINER_MODULE = "trainer.model"
TRAIN_DIR = os.getcwd() + '/trainer'
JOB_DIR = BUCKET_URL + '/output'
REGION = "us-central1"

Cada job de treinamento no AI Platform precisa ter um nome exclusivo. Execute o comando a seguir para definir uma variável para o nome do job usando um carimbo de data/hora:

timestamp = str(datetime.datetime.now().time())
JOB_NAME = 'caip_training_' + str(int(time.time()))

Etapa 2: iniciar o job de treinamento

Vamos criar nosso job de treinamento usando a gcloud, a CLI do Google Cloud. Podemos executar esse comando diretamente no notebook, referenciando as variáveis definidas acima:

!gcloud ai-platform jobs submit training $JOB_NAME \
        --scale-tier basic \
        --package-path $TRAIN_DIR \
        --module-name $MAIN_TRAINER_MODULE \
        --job-dir $JOB_DIR \
        --region $REGION \
        --runtime-version 2.1 \
        --python-version 3.7 \
        --config config.yaml

Se o job foi criado corretamente, acesse a seção Jobs do console do AI Platform para monitorar os registros.

Etapa 3: monitorar seu job

Na seção "Jobs" do console, clique no job que você acabou de iniciar para ver os detalhes:

À medida que a primeira rodada de testes começa, você pode conferir os valores de hiperparâmetros selecionados para cada teste:

À medida que os testes forem concluídos, o valor resultante da métrica de otimização (neste caso, val_mse) será registrado aqui. O job leva de 15 a 20 minutos para ser executado, e o painel fica assim quando ele termina (os valores exatos variam):

Para depurar possíveis problemas e monitorar seu job com mais detalhes, clique em Ver registros na página de detalhes do job:

Todas as instruções print() no código de treinamento de modelo vão aparecer aqui. Se você estiver enfrentando problemas, adicione mais instruções de impressão e inicie um novo job de treinamento.

Depois que o job de treinamento for concluído, encontre os hiperparâmetros que geraram o menor val_mse. Você pode usar esses dados para treinar e exportar uma versão final do modelo ou como orientação para iniciar outro job de treinamento com mais testes de ajuste de hiperparâmetros.

Está criando seu modelo com outro framework? Mostramos como treinar um modelo do PyTorch aqui, mas o processo de treinamento e ajuste de modelos criados com qualquer outro framework no AI Platform é exatamente o mesmo. Basta substituir o código do modelo.

7. Limpeza

Se você quiser continuar usando esse notebook, recomendamos que o desative quando não estiver em uso. Na interface de Notebooks no console do Cloud, selecione o notebook e clique em Parar:

Se quiser excluir todos os recursos criados neste laboratório, exclua a instância do notebook em vez de interrompê-la.

No menu de navegação do console do Cloud, acesse "Storage" e exclua os dois buckets criados para armazenar os recursos do modelo.