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

1. Descripción general

En este lab, seguirás un flujo de trabajo completo de entrenamiento de AA en Google Cloud con PyTorch para compilar tu modelo. En un entorno de Cloud AI Platform Notebooks, 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. Configura 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 GPUs):

Usa las opciones predeterminadas o asígnale un nombre personalizado si lo deseas y, luego, 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 puedes comenzar.

Paso 5: Importa paquetes de Python

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

En la primera celda del notebook, agrega las siguientes importaciones y ejecuta la celda. Para ejecutarlo, presiona el botón de flecha hacia la derecha en el menú superior o presiona Comando + 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 ejecutamos 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 que nuestro código de entrenamiento esté empaquetado de forma local en nuestra instancia de Notebooks y un bucket de Cloud Storage para almacenar los recursos de nuestro trabajo. Primero, crearemos un bucket de almacenamiento. Puedes omitir este paso si ya tienes una.

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 valores a continuación con el nombre de tu proyecto de Google Cloud y el nombre del bucket de Cloud Storage que deseas 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 podemos crear un bucket de almacenamiento, al que haremos referencia cuando iniciemos el trabajo de entrenamiento.

Ejecuta este comando de 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, deberemos configurar nuestro código como un paquete de Python. Esto consiste en un archivo setup.py en nuestro directorio raíz que especifica las dependencias de paquetes externos, un subdirectorio con el nombre de nuestro 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 dentro de él. Python usa este archivo para reconocer que se trata de un paquete:

!mkdir trainer
!touch trainer/__init__.py

Ahora sí podemos comenzar a crear nuestro trabajo de entrenamiento.

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

El enfoque de este lab se centra en las herramientas para entrenar modelos, pero echemos un vistazo rápido al conjunto de datos que usaremos para entrenar nuestro modelo y comprenderlo. Usaremos el conjunto de datos de natalidad disponible en BigQuery. Contiene datos de nacimientos de EE.UU. de varias décadas. Usaremos algunas columnas del conjunto de datos para predecir el peso de un bebé al nacer. El conjunto de datos original es bastante grande, por lo que usaremos un subconjunto que pusimos a tu disposición en un bucket de Cloud Storage.

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

Descarguemos la versión del conjunto de datos que pusimos a tu disposición en Cloud Storage en un DataFrame de Pandas y obtengamos una vista previa.

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

Este conjunto de datos tiene 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 sexo del bebé representado como un valor booleano.

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

Escribiremos nuestro script de entrenamiento en un archivo llamado model.py dentro del subdirectorio trainer/ que creamos antes. Nuestro trabajo de entrenamiento se ejecutará en AI Platform Training y también usará el servicio de ajuste de hiperparámetros de AI Platform para encontrar los hiperparámetros óptimos para nuestro modelo con la optimización bayesiana.

Paso 1: Crea la secuencia de comandos de entrenamiento

Primero, creemos el archivo de Python con nuestro script de entrenamiento. Luego, analizaremos lo que sucede en él. Si ejecutas este comando %%writefile, se escribirá el código del modelo en un archivo de 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()

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

• get_args(): Analiza los argumentos de la línea de comandos que pasaremos cuando creemos nuestro trabajo de entrenamiento, junto con los hiperparámetros que queremos que AI Platform optimice. En este ejemplo, nuestra lista de argumentos solo incluye los hiperparámetros que optimizaremos: la tasa de aprendizaje, el momentum y la cantidad de neuronas de la capa oculta de nuestro modelo.
• train_model(): Aquí descargamos los datos en un DataFrame de Pandas, los normalizamos, los convertimos en tensores de PyTorch y, luego, definimos nuestro modelo. Para compilar nuestro modelo, usamos la API de nn.Sequential de PyTorch, que nos permite definir nuestro 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 de forma rígida el tamaño de la capa oculta de nuestro modelo, lo convertimos en un hiperparámetro que AI Platform ajustará por nosotros. Hablaremos más sobre este tema en la siguiente sección.

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

En lugar de probar manualmente diferentes valores de hiperparámetros y volver a entrenar nuestro modelo cada vez, usaremos el servicio de optimización de hiperparámetros de Cloud AI Platform. 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 especifiquemos.

En el ajuste de hiperparámetros, una sola 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 para optimizar los hiperparámetros que seleccione para las pruebas futuras. Para configurar el ajuste de hiperparámetros, debemos pasar un archivo de configuración cuando iniciemos nuestro trabajo de entrenamiento con algunos datos sobre cada uno de los hiperparámetros 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 aumentará el valor en las diferentes pruebas.

Al comienzo del trabajo, también especificamos la métrica para la que realizamos la optimización. Ten en cuenta que, al final de nuestra 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 generen 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 informarla cuando llamamos a report_hyperparameter_tuning_metric() al final de una prueba.

Nota: Si bien no ejecutamos nuestro trabajo de entrenamiento directamente en nuestra 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 por valores de hiperparámetros codificados. Como ya creaste una instancia de PyTorch, puedes continuar.

6. Ejecuta un trabajo de entrenamiento en AI Platform

En esta sección, iniciaremos nuestro trabajo de entrenamiento del modelo 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 que se muestra 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 con 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, la CLI de Google Cloud. Podemos ejecutar este comando directamente en nuestro notebook, haciendo referencia a las variables que definimos anteriormente:

!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

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

A medida que 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, se registrará aquí el valor resultante de tu métrica de optimización (en este caso, val_mse). El trabajo debería tardar entre 15 y 20 minutos en ejecutarse, y el panel se verá de la siguiente manera cuando finalice (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í aparecerá cada instrucción print() en el código de entrenamiento del modelo. Si tienes problemas, intenta agregar más instrucciones de impresión y comienza un nuevo trabajo de entrenamiento.

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

¿Compilas tu modelo con otro framework? Aquí te mostramos cómo entrenar un modelo de PyTorch, pero el proceso para entrenar y ajustar 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 usando este notebook, te recomendamos que lo desactives cuando no lo uses. 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, simplemente borra la instancia del notebook en lugar de detenerla.

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