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

1. Visão geral

Neste laboratório, você analisará um fluxo de trabalho completo de treinamento de ML no Google Cloud usando o PyTorch para criar seu modelo. Em um ambiente de Notebooks do AI Platform do Cloud, você 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 PyTorch
• Treine seu modelo com o 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", caso essa opção ainda não esteja ativada.

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 de Notebooks no AI Platform

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

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

Em seguida, abra uma instância do Notebook do Python 3 no acesso rápido:

Você já pode começar.

Etapa 5: importar pacotes do 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 abaixo e execute a célula. Para executá-lo, 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

Observe que não estamos importando o PyTorch aqui. Isso ocorre porque estamos executando o job de treinamento no AI Platform Training, não na nossa instância do Notebook.

3. Criar um pacote para o job de treinamento

Para executar nosso job de treinamento no AI Platform Training, vamos precisar do nosso código de treinamento empacotado localmente na nossa instância de Notebooks e de um bucket do Cloud Storage para armazenar recursos para o job. Primeiro, criaremos um bucket de armazenamento. Se você já tiver uma, pule esta etapa.

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

Primeiro, defina algumas variáveis de ambiente que serão usadas no restante do codelab. Preencha os valores abaixo com o nome do seu projeto do Google Cloud e 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 estamos prontos para criar um bucket de armazenamento, que vamos apontar quando iniciarmos nosso job de treinamento.

Execute este comando gsutil no seu notebook para criar um bucket:

!gsutil mb $BUCKET_URL

Etapa 2: criar os arquivos iniciais do nosso pacote Python

Para executar um job de treinamento no AI Platform, será preciso configurar nosso código como um pacote do Python. Ele consiste em um arquivo setup.py no diretório raiz que especifica todas as dependências de pacotes externas, um subdiretório com o nome do pacote (chamado aqui de trainer/) e um arquivo __init__.py vazio nesse subdiretório.

Primeiro, vamos criar nosso arquivo setup.py. Estamos usando os comandos mágicos %%writefile do iPython para salvar o arquivo na nossa instância. Especificamos três bibliotecas externas que vamos usar 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.'
)

Agora vamos criar o 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 estamos prontos para começar a criar nosso job de treinamento.

4. Visualizar o conjunto de dados

O foco deste laboratório são as ferramentas para treinar modelos. Mas vamos analisar rapidamente o conjunto de dados que vamos usar no treinamento do 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. Usaremos 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 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 disponibilizadas 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()

Este conjunto de dados tem pouco menos de 100.000 linhas. Usaremos cinco atributos para prever o peso de um bebê ao nascer: idade da mãe e do pai, semanas de gestação, ganho de peso da mãe em libras e gênero do bebê representado como um booleano.

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

Escreveremos nosso script de treinamento em um arquivo chamado model.py dentro do subdiretório trainer/ criado anteriormente. Nosso job de treinamento será executado no AI Platform Training e 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, analisaremos o que está acontecendo nele. A execução deste comando %%writefile 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 trabalho 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 serão transmitidos 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 do modelo, o momentum e o número de neurônios na camada escondida.
• train_model(): aqui fazemos o download dos dados em um DataFrame do Pandas, normalizamos, convertemos em tensores PyTorch e definimos nosso modelo. Para criar o modelo, estamos usando a API PyTorch nn.Sequential, 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))

Em vez de fixar o tamanho da camada escondida do modelo no código, estamos tornando esse um hiperparâmetro que o AI Platform vai ajustar. Falaremos mais sobre isso na próxima seção.

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

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

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 executarmos, a AI Platform usará os resultados dos testes concluídos para otimizar os hiperparâmetros selecionados para futuros. Para configurar o ajuste de hiperparâmetros, precisamos passar um arquivo de configuração quando iniciarmos nosso job de treinamento com alguns dados em 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 gostaríamos de pesquisar e a escala para aumentar o valor em diferentes testes.

No início do job, também especificamos a métrica para a qual estamos otimizando. No final da função train_model() acima, informamos essa métrica à AI Platform sempre que um teste é concluído. Aqui estamos minimizando o erro quadrático médio do nosso modelo. Por isso, queremos usar os hiperparâmetros que resultam no menor erro quadrático médio do 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: não estamos executando o job de treinamento diretamente na instância do notebook, mas você pode copiar o código de treinamento do modelo em uma célula nova e testá-lo no notebook. Basta substituir todas as referências a args por valores de hiperparâmetros fixados no código. Como você criou uma instância do PyTorch, já pode começar.

6. execute um job de treinamento no AI Platform

Nesta seção, iniciaremos nosso job de treinamento de modelo com o ajuste de hiperparâmetros no AI Platform.

Etapa 1: definir algumas variáveis de ambiente

Primeiro, defina algumas variáveis de ambiente que serão usadas para iniciar nosso job de treinamento. Se você quiser executar seu 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 seguinte comando 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 em nosso notebook, fazendo referência às 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 tiver sido criado corretamente, acesse a seção Jobs do console do AI Platform para monitorar os registros.

Etapa 3: monitorar o job

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

No início da primeira rodada de testes, você poderá conferir os valores de hiperparâmetros selecionados para cada teste:

Quando os testes forem concluídos, o valor resultante da métrica de otimização (neste caso, val_mse) será registrado aqui. O job deve levar de 15 a 20 minutos para ser executado, e o painel ficará mais ou menos assim quando o job for concluído (os valores exatos variarão):

Para depurar possíveis problemas e monitorar o 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 do modelo serão exibidas aqui. Se você tiver problemas, tente adicionar mais instruções de impressão e iniciar um novo trabalho de treinamento.

Quando o job de treinamento for concluído, encontre os hiperparâmetros que produziram o menor val_mse. É possível 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 outros testes de ajuste de hiperparâmetros.

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

7. Limpeza

Se você quiser continuar usando este notebook, é recomendado que você o desative quando não estiver em uso. A partir da IU de Notebooks no Console do Cloud, selecione o notebook e depois clique em Parar:

Se 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 do console do Cloud, acesse "Storage" e exclua os dois buckets que você criou para armazenar os recursos do modelo.