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

1. Visão geral

Neste laboratório, você vai passar por um fluxo de trabalho completo de treinamento de ML no Google Cloud, usando o PyTorch para criar seu modelo. Em um ambiente de Notebooks da IA do Google Cloud Platform, 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 de notebooks do AI Platform
• Criar um modelo do PyTorch
• Treinar seu 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 as instruções aqui.

Etapa 1: ativar a API Models da IA do Google Cloud

Navegue até a seção AI Platform Models do console do Cloud e clique em "Ativar" se ela ainda não estiver ativada.

Etapa 2: ativar a API Compute Engine

Navegue até Compute Engine e selecione Ativar se ele ainda não estiver ativado. Você vai precisar disso para criar sua instância de notebook.

Etapa 3: criar uma instância de notebooks do AI Platform

Navegue até 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 no menu inicial:

Tudo pronto para começar!

Etapa 5: importar pacotes do Python

Para o restante deste codelab, execute todos os snippets de código no notebook do Jupyter.

Na primeira célula do notebook, adicione as seguintes importações 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 ocorre 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 que o código de treinamento seja empacotado localmente na instância de Notebooks e um bucket do Cloud Storage para armazenar recursos do job. Primeiro, vamos criar um bucket de armazenamento. Você pode pular esta etapa se já tiver um.

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

Primeiro, vamos definir algumas variáveis de ambiente que serão usadas no restante do codelab. Preencha os valores abaixo com o nome do seu projeto na nuvem do Google 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 já podemos 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 do Python

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

Primeiro, vamos escrever nosso arquivo setup.py. Estamos usando as mágicas %%writefile do iPython para salvar o arquivo na instância. Aqui, especificamos três bibliotecas externas que serão usadas no 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 nosso diretório trainer/ e o arquivo init.py vazio nele. O Python usa esse arquivo para reconhecer que se trata de um pacote:

!mkdir trainer
!touch trainer/__init__.py

Agora já podemos começar a criar nosso job de treinamento.

4. Visualizar o conjunto de dados

O foco deste laboratório está nas ferramentas para treinar modelos, mas vamos dar uma olhada rápida no conjunto de dados que vamos usar para treinar nosso modelo. Vamos usar o conjunto de dados de natalidade disponível no BigQuery. Ele contém dados de nascimento dos EUA ao longo de várias décadas. Vamos usar algumas colunas do conjunto de dados para prever o peso de nascimento de um bebê. O conjunto de dados original é bastante 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 fazer o download da versão do conjunto de dados que disponibilizamos para você no Cloud Storage para um DataFrame do Pandas e visualizá-lo.

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 o 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. Em seguida, vamos analisar o que está acontecendo nele. A execução desse 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 nosso job de treinamento, juntamente com os 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 fazemos o download dos dados para um DataFrame do Pandas, normalizamos, convertemos em tensores do PyTorch e definimos nosso modelo. Para criar nosso modelo, estamos usando a API nn.Sequential do PyTorch, que nos permite definir nosso 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 tornando isso um hiperparâmetro que o AI Platform vai ajustar para nós. Mais informações sobre isso na próxima seção.

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

Em vez de tentar manualmente valores de hiperparâmetros diferentes e treinar nosso modelo novamente a cada vez, vamos usar o serviço de otimização de hiperparâmetros do AI Platform da IA 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 usará 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 esse 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, reportamos essa métrica ao AI Platform sempre que um teste é concluído. Aqui, estamos minimizando o erro quadrático médio do modelo e, portanto, queremos usar os hiperparâmetros que resultam no erro quadrático médio mais baixo para nosso modelo. O nome dessa métrica (val_mse) corresponde ao nome que usamos para informá-la quando chamamos report_hyperparameter_tuning_metric() no 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 testá-lo no notebook. Basta substituir todas as referências a args por valores de hiperparâmetros codificados. Como você 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 no AI Platform.

Etapa 1: definir algumas variáveis de ambiente

Primeiro, vamos definir algumas variáveis de ambiente que serão usadas para iniciar nosso 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 o job

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

À medida que a primeira rodada de testes é iniciada, você poderá ver os valores de hiperparâmetros selecionados para cada teste:

À medida que os testes são 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 parecido com este quando o job termina (os valores exatos variam):

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

Todas as instruções print() no código de treinamento de modelo serão mostradas aqui. Se você estiver enfrentando problemas, tente adicionar mais instruções de impressão e iniciar um novo job de treinamento.

Depois que o job de treinamento for concluído, encontre os hiperparâmetros que produziram o val_mse mais baixo. Você pode usá-los para treinar e exportar uma versão final do modelo ou usá-los como orientação para iniciar outro job de treinamento com mais testes de ajuste de hiperparâmetros.

Criando seu modelo com outra estrutura? Mostramos como treinar um modelo do PyTorch aqui, mas o processo de treinamento e ajuste de modelos criados com qualquer outra estrutura no AI Platform é exatamente o mesmo! Basta substituir o código do modelo.

7. Revisão dos dados

Se você quiser continuar usando este notebook, é recomendável que você o desative quando não estiver em uso. Na interface de notebooks no console do Cloud, selecione o notebook e clique em Parar:

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

Usando o menu de navegação no console do Cloud, navegue até o Storage e exclua os dois buckets criados para armazenar os recursos do modelo.