Entrenamiento y ajuste de hiperparámetros de un modelo de PyTorch en AI Platform de Cloud

1. Descripción general

En este lab, recorrerás un flujo de trabajo completo de entrenamiento de AA en Google Cloud con PyTorch para compilar tu modelo. Desde un entorno de AI Platform Notebooks de Cloud, aprenderás a empaquetar tu trabajo de entrenamiento para ejecutarlo en AI Platform Training con ajuste de hiperparámetros.

Qué aprenderá

Aprenderás a hacer lo siguiente:

• Crea una instancia de AI Platform Notebooks
• Crea un modelo de PyTorch
• Entrena tu modelo con el ajuste de hiperparámetros en AI Platform Training

El costo total de la ejecución de este lab en Google Cloud es de aproximadamente $1.

2. Cómo configurar tu entorno

Para ejecutar este codelab, necesitarás un proyecto de Google Cloud Platform que tenga habilitada la facturación. Para crear un proyecto, sigue estas instrucciones.

Paso 1: Habilita la API de Cloud AI Platform Models

Navega a la sección Modelos de AI Platform de la consola de Cloud y haz clic en Habilitar si aún no está habilitada.

Paso 2: Habilita la API de Compute Engine

Ve a Compute Engine y selecciona Habilitar si aún no está habilitada. La necesitarás para crear la instancia de notebook.

Paso 3: Crea una instancia de AI Platform Notebooks

Navega a la sección AI Platform Notebooks de la consola de Cloud y haz clic en Instancia nueva. Luego, selecciona el tipo de instancia de PyTorch más reciente (sin GPU):

Usa las opciones predeterminadas o asígnale un nombre personalizado, y haz clic en Crear. Una vez que se crea la instancia, selecciona Abrir JupyterLab:

A continuación, abre una instancia de notebook de Python 3 desde el selector:

¡Ya estás listo para empezar!

Paso 5: Importa paquetes de Python

Para el resto de este codelab, ejecuta todos los fragmentos de código de tu notebook de Jupyter.

En la primera celda de tu notebook, agrega las siguientes importaciones y ejecuta la celda. Para ejecutarla, presiona el botón de la flecha hacia la derecha en el menú superior o presiona Command + Intro:

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

Notarás que no importamos PyTorch aquí. Esto se debe a que estamos ejecutando el trabajo de entrenamiento en AI Platform Training, no desde nuestra instancia de notebook.

3. Crea un paquete para el trabajo de entrenamiento

Para ejecutar nuestro trabajo de entrenamiento en AI Platform Training, necesitaremos el código de entrenamiento empaquetado de forma local en nuestra instancia de Notebooks y un bucket de Cloud Storage para almacenar elementos para nuestro trabajo. Primero, crearemos un bucket de almacenamiento. Puedes omitir este paso si ya tienes uno.

Paso 1: Crea un bucket de Cloud Storage para nuestro modelo

Primero, definamos algunas variables de entorno que usaremos durante el resto del codelab. Completa los siguientes valores con el nombre de tu proyecto de Google Cloud y el nombre del bucket de Cloud Storage que quieres crear (debe ser único a nivel global):

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

Ahora estamos listos para crear un bucket de almacenamiento, al que señalaremos cuando iniciemos nuestro trabajo de entrenamiento.

Ejecuta este comando gsutil desde tu notebook para crear un bucket:

!gsutil mb $BUCKET_URL

Paso 2: Crea los archivos iniciales para nuestro paquete de Python

Para ejecutar un trabajo de entrenamiento en AI Platform, tendremos que configurar el código como un paquete de Python. Este archivo consta de un archivo setup.py en nuestro directorio raíz que especifica las dependencias de paquetes externas, un subdirectorio con el nombre del paquete (aquí lo llamaremos trainer/) y un archivo __init__.py vacío dentro de este subdirectorio.

Primero, escribamos nuestro archivo setup.py. Usamos los comandos mágicos %%writefile de iPython para guardar el archivo en nuestra instancia. Aquí, especificamos 3 bibliotecas externas que usaremos en nuestro código de entrenamiento: PyTorch, scikit-learn y 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.'
)

A continuación, crearemos nuestro directorio trainer/ y el archivo init.py vacío que contiene. Python usa este archivo para reconocer que se trata de un paquete:

!mkdir trainer
!touch trainer/__init__.py

Ahora estamos listos para comenzar a crear nuestro trabajo de entrenamiento.

4. Obtén una vista previa del conjunto de datos

El enfoque de este lab está en las herramientas para entrenar modelos aquí, pero repasemos rápidamente el conjunto de datos que usaremos para entrenar nuestro modelo de modo que lo entienda. Utilizaremos el conjunto de datos de natalidad disponible en BigQuery. Contiene datos de nacimientos en EE.UU. durante varias décadas. Utilizaremos algunas columnas del conjunto de datos para predecir el peso de un bebé al nacer. El conjunto de datos original es bastante grande y usaremos un subconjunto de él que pusimos a tu disposición en un bucket de Cloud Storage.

Paso 1: Descarga el conjunto de datos de natalidad de BigQuery

Descargue a un DataFrame de Pandas la versión del conjunto de datos que pusimos a su disposición en Cloud Storage y obtenga una vista previa.

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

Este conjunto de datos tiene un poco menos de 100,000 filas. Usaremos 5 atributos para predecir el peso de un bebé al nacer: la edad de la madre y el padre, las semanas de gestación, el aumento de peso de la madre en libras y el género del bebé representado como un valor booleano.

5. Define el trabajo de entrenamiento con ajuste de hiperparámetros

Escribiremos nuestra secuencia de comandos de entrenamiento en un archivo llamado model.py dentro del subdirectorio o entrenador que creamos anteriormente. Nuestro trabajo de entrenamiento se ejecutará en AI Platform Training y también usará el servicio de ajuste de hiperparámetros de AI Platform a fin de encontrar los hiperparámetros óptimos para nuestro modelo mediante la optimización bayesiana.

Paso 1: Crea la secuencia de comandos de entrenamiento

Primero, creemos el archivo de Python con nuestra secuencia de comandos de entrenamiento. Luego, desglosaremos lo que sucede en él. Si ejecutas este comando %%writefile, se escribirá el código del modelo en un archivo local de Python:

%%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()

El trabajo de entrenamiento consta de dos funciones en las que se realiza la mayor parte del trabajo.

• get_args(): Esto analiza los argumentos de la línea de comandos que pasaremos cuando creemos el trabajo de entrenamiento, junto con los hiperparámetros que queremos que AI Platform optimice. En este ejemplo, nuestra lista de argumentos incluye solo los hiperparámetros que optimizaremos: la tasa de aprendizaje, el momentum y la cantidad de neuronas de nuestro modelo en la capa oculta.
• train_model(): Aquí descargamos los datos en un DataFrame de Pandas, los normalizamos, los convertimos en tensores de PyTorch y, luego, definimos el modelo. Para compilar nuestro modelo, usamos la API de nn.Sequential de PyTorch, que nos permite definir el modelo como una pila de capas:

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

Ten en cuenta que, en lugar de codificar el tamaño de la capa oculta de nuestro modelo, lo convertiremos en un hiperparámetro que AI Platform ajustará por nosotros. Más sobre este tema en la próxima sección.

Paso 2: Usa el servicio de ajuste de hiperparámetros de AI Platform

En lugar de probar diferentes valores de hiperparámetros y volver a entrenar nuestro modelo de forma manual cada vez, usaremos el servicio de optimización de hiperparámetros de AI Platform de Cloud. Si configuramos nuestro trabajo de entrenamiento con argumentos de hiperparámetros, AI Platform usará la optimización bayesiana para encontrar los valores ideales para los hiperparámetros que especificamos.

En el ajuste de hiperparámetros, una única prueba consiste en una ejecución de entrenamiento de nuestro modelo con una combinación específica de valores de hiperparámetros. Según la cantidad de pruebas que ejecutemos, AI Platform usará los resultados de las pruebas completadas a fin de optimizar los hiperparámetros que seleccione para los futuros. Para configurar el ajuste de hiperparámetros, debemos pasar un archivo de configuración cuando iniciamos nuestro trabajo de entrenamiento con algunos datos sobre cada hiperparámetro que optimizamos.

A continuación, crea ese archivo de configuración de forma local:

%%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 el tipo, el rango de valores que queremos buscar y la escala en la que se debe aumentar el valor en diferentes pruebas.

Al comienzo del trabajo, también especificamos la métrica para la que estamos optimizando. Ten en cuenta que, al final de la función train_model() anterior, informamos esta métrica a AI Platform cada vez que se completa una prueba. Aquí minimizamos el error cuadrático medio de nuestro modelo, por lo que queremos usar los hiperparámetros que dan como resultado el error cuadrático medio más bajo para nuestro modelo. El nombre de esta métrica (val_mse) coincide con el nombre que usamos para registrarla cuando llamamos a report_hyperparameter_tuning_metric() al final de una prueba.

Nota: Si bien no estamos ejecutando nuestro trabajo de entrenamiento directamente en la instancia de notebook, puedes copiar el código de entrenamiento del modelo en una celda nueva y probarlo en tu notebook. Solo reemplaza las referencias a args con valores de hiperparámetros codificados. Como creaste una instancia de PyTorch, está todo listo.

6. Ejecuta un trabajo de entrenamiento en AI Platform

En esta sección, iniciaremos el trabajo de entrenamiento de modelos con el ajuste de hiperparámetros en AI Platform.

Paso 1: Define algunas variables de entorno

Primero, definamos algunas variables de entorno que usaremos para iniciar nuestro trabajo de entrenamiento. Si quieres ejecutar tu trabajo en una región diferente, actualiza la variable REGION a continuación:

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

Cada trabajo de entrenamiento en AI Platform debe tener un nombre único. Ejecuta el siguiente comando para definir una variable para el nombre de tu trabajo usando una marca de tiempo:

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

Paso 2: Inicia el trabajo de entrenamiento

Crearemos nuestro trabajo de entrenamiento con gcloud, Google Cloud CLI. Podemos ejecutar este comando directamente en nuestro notebook, haciendo referencia a las variables que definimos antes:

!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

Si tu trabajo se creó correctamente, ve a la sección Trabajos de la consola de AI Platform para supervisar los registros.

Paso 3: Supervisa tu trabajo

Cuando estés en la sección Trabajos de la consola, haz clic en el trabajo que acabas de empezar para ver los detalles:

Cuando comience la primera ronda de pruebas, podrás ver los valores de hiperparámetros seleccionados para cada prueba:

A medida que se completen las pruebas, el valor resultante de tu métrica de optimización (en este caso, val_mse) se registrará aquí. El trabajo debería tardar entre 15 y 20 minutos en ejecutarse, y el panel se verá así cuando el trabajo haya finalizado (los valores exactos variarán):

Para depurar posibles problemas y supervisar tu trabajo con más detalle, haz clic en Ver registros en la página de detalles del trabajo:

Aquí se mostrará cada sentencia print() en el código de entrenamiento del modelo. Si tienes problemas, intenta agregar más declaraciones de impresión y comenzar un nuevo trabajo de entrenamiento.

Una vez que se complete el trabajo de entrenamiento, busca los hiperparámetros que produjeron el valor más bajo de val_mse. Puedes usarlos para entrenar y exportar una versión final de tu modelo o usarlos como guía para iniciar otro trabajo de entrenamiento con pruebas adicionales de ajuste de hiperparámetros.

¿Estás compilando tu modelo con otro framework? Aquí te mostramos cómo entrenar un modelo de PyTorch, pero el proceso de entrenamiento y ajuste de modelos creados con cualquier otro framework en AI Platform es exactamente el mismo. Solo reemplaza el código del modelo.

7. Limpieza

Si quieres seguir utilizando este bloc de notas, te recomendamos que lo desactives cuando no lo utilices. En la IU de Notebooks de la consola de Cloud, selecciona el notebook y, luego, haz clic en Detener:

Si quieres borrar todos los recursos que creaste en este lab, solo borra la instancia de notebook en lugar de detenerla.

En el menú de navegación de la consola de Cloud, navega a Almacenamiento y borra los buckets que creaste para almacenar tus recursos del modelo.