Créer un script d'entraînement

Pour créer un modèle personnalisé, vous avez besoin d'un script d'entraînement Python qui crée et entraîne le modèle personnalisé. Vous initialisez votre job d'entraînement avec le script d'entraînement Python, puis appelez la méthode run du job d'entraînement pour exécuter le script.

Dans cette rubrique, vous allez créer le script d'entraînement, puis spécifier des arguments de commande pour votre script d'entraînement.

Créer un script d'entraînement

Dans cette section, vous allez créer un script d'entraînement. Ce script est un nouveau fichier nommé task.py dans votre environnement de notebook. Plus loin dans ce tutoriel, vous transmettrez ce script au constructeur aiplatform.CustomTrainingJob. Lorsque le script s'exécute, il effectue les opérations suivantes :

• Il charge les données dans l'ensemble de données BigQuery que vous avez créé.

• Il exploite l'API TensorFlow Keras pour créer, compiler et entraîner votre modèle.

• Il spécifie le nombre d'époques et la taille de lot à utiliser lorsque la méthode Keras Model.fit est appelée.

• Il spécifie où enregistrer les artefacts de modèle à l'aide de la variable d'environnement AIP_MODEL_DIR. AIP_MODEL_DIR est défini par Vertex AI et contient l'URI d'un répertoire pour l'enregistrement des artefacts de modèle. Pour en savoir plus, consultez la section Variables d'environnement pour les répertoires Cloud Storage spéciaux.

• Il exporte un modèle SavedModel TensorFlow vers le répertoire de modèle. Pour en savoir plus, consultez la section Utiliser le format SavedModel sur le site Web de TensorFlow.

Pour créer votre script d'entraînement, exécutez le code suivant dans votre notebook :

%%writefile task.py

import argparse
import numpy as np
import os

import pandas as pd
import tensorflow as tf

from google.cloud import bigquery
from google.cloud import storage

# Read environmental variables
training_data_uri = os.getenv("AIP_TRAINING_DATA_URI")
validation_data_uri = os.getenv("AIP_VALIDATION_DATA_URI")
test_data_uri = os.getenv("AIP_TEST_DATA_URI")

# Read args
parser = argparse.ArgumentParser()
parser.add_argument('--label_column', required=True, type=str)
parser.add_argument('--epochs', default=10, type=int)
parser.add_argument('--batch_size', default=10, type=int)
args = parser.parse_args()

# Set up training variables
LABEL_COLUMN = args.label_column

# See https://cloud.google.com/vertex-ai/docs/workbench/managed/executor#explicit-project-selection for issues regarding permissions.
PROJECT_NUMBER = os.environ["CLOUD_ML_PROJECT_ID"]
bq_client = bigquery.Client(project=PROJECT_NUMBER)


# Download a table
def download_table(bq_table_uri: str):
    # Remove bq:// prefix if present
    prefix = "bq://"
    if bq_table_uri.startswith(prefix):
        bq_table_uri = bq_table_uri[len(prefix) :]
        
    # Download the BigQuery table as a dataframe
    # This requires the "BigQuery Read Session User" role on the custom training service account.
    table = bq_client.get_table(bq_table_uri)
    return bq_client.list_rows(table).to_dataframe()

# Download dataset splits
df_train = download_table(training_data_uri)
df_validation = download_table(validation_data_uri)
df_test = download_table(test_data_uri)

def convert_dataframe_to_dataset(
    df_train: pd.DataFrame,
    df_validation: pd.DataFrame,
):
    df_train_x, df_train_y = df_train, df_train.pop(LABEL_COLUMN)
    df_validation_x, df_validation_y = df_validation, df_validation.pop(LABEL_COLUMN)

    y_train = tf.convert_to_tensor(np.asarray(df_train_y).astype("float32"))
    y_validation = tf.convert_to_tensor(np.asarray(df_validation_y).astype("float32"))

    # Convert to numpy representation
    x_train = tf.convert_to_tensor(np.asarray(df_train_x).astype("float32"))
    x_test = tf.convert_to_tensor(np.asarray(df_validation_x).astype("float32"))

    # Convert to one-hot representation
    num_species = len(df_train_y.unique())
    y_train = tf.keras.utils.to_categorical(y_train, num_classes=num_species)
    y_validation = tf.keras.utils.to_categorical(y_validation, num_classes=num_species)

    dataset_train = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    dataset_validation = tf.data.Dataset.from_tensor_slices((x_test, y_validation))
    return (dataset_train, dataset_validation)

# Create datasets
dataset_train, dataset_validation = convert_dataframe_to_dataset(df_train, df_validation)

# Shuffle train set
dataset_train = dataset_train.shuffle(len(df_train))

def create_model(num_features):
    # Create model
    Dense = tf.keras.layers.Dense
    model = tf.keras.Sequential(
        [
            Dense(
                100,
                activation=tf.nn.relu,
                kernel_initializer="uniform",
                input_dim=num_features,
            ),
            Dense(75, activation=tf.nn.relu),
            Dense(50, activation=tf.nn.relu),            
            Dense(25, activation=tf.nn.relu),
            Dense(3, activation=tf.nn.softmax),
        ]
    )
    
    # Compile Keras model
    optimizer = tf.keras.optimizers.RMSprop(lr=0.001)
    model.compile(
        loss="categorical_crossentropy", metrics=["accuracy"], optimizer=optimizer
    )
    
    return model

# Create the model
model = create_model(num_features=dataset_train._flat_shapes[0].dims[0].value)

# Set up datasets
dataset_train = dataset_train.batch(args.batch_size)
dataset_validation = dataset_validation.batch(args.batch_size)

# Train the model
model.fit(dataset_train, epochs=args.epochs, validation_data=dataset_validation)

tf.saved_model.save(model, os.getenv("AIP_MODEL_DIR"))

Une fois créé, le script apparaît dans le dossier racine de votre notebook :

Définir des arguments pour votre script d'entraînement

Vous transmettez les arguments de ligne de commande suivants à votre script d'entraînement :

• label_column : identifie la colonne de vos données qui contient ce que vous souhaitez prédire. Dans le cas présent, cette colonne est species. Vous avez défini cette option dans une variable nommée LABEL_COLUMN lors du traitement de vos données. Pour en savoir plus, consultez la section Télécharger, prétraiter et diviser les données.

• epochs : nombre d'époques utilisées lors de l'entraînement de votre modèle. Une époque est une itération sur les données lors de l'entraînement de votre modèle. Ce tutoriel utilise 20 époques.

• batch_size : nombre d'échantillons traités avant la mise à jour de votre modèle. Ce tutoriel utilise une taille de lot de 10.

Pour définir les arguments transmis à votre script, exécutez le code suivant :

JOB_NAME = "custom_job_unique"

EPOCHS = 20
BATCH_SIZE = 10

CMDARGS = [
    "--label_column=" + LABEL_COLUMN,
    "--epochs=" + str(EPOCHS),
    "--batch_size=" + str(BATCH_SIZE),
]