Crea uno script di addestramento

Per creare un modello personalizzato, devi avere uno script di addestramento Python che crei e addestri il modello personalizzato. Inizializzi il job di addestramento con lo script di addestramento Python, quindi richiami il metodo run del job di addestramento per eseguire lo script.

In questo argomento, crei lo script di addestramento, quindi specifichi gli argomenti del comando per lo script di addestramento.

Crea uno script di addestramento

In questa sezione viene creato uno script di addestramento. Questo script è un nuovo file nell'ambiente del tuo notebook denominato task.py. Più avanti in questo tutorial, passerai questo script al costruttore aiplatform.CustomTrainingJob. Quando viene eseguito, lo script esegue le seguenti operazioni:

  • Carica i dati nel set di dati BigQuery che hai creato.

  • Utilizza l'API TensorFlow Keras per compilare, creare e addestrare il modello.

  • Specifica il numero di epoche e la dimensione del batch da utilizzare quando viene invocato il metodo Keras Model.fit.

  • Specifica dove salvare gli elementi del modello utilizzando la variabile di ambiente AIP_MODEL_DIR. AIP_MODEL_DIR viene impostato da Vertex AI e contiene l'URI di una directory per il salvataggio degli elementi del modello. Per ulteriori informazioni, consulta Variabili di ambiente per le directory Cloud Storage speciali.

  • Esporta un SavedModel di TensorFlow nella directory del modello. Per ulteriori informazioni, consulta la sezione Utilizzare il formato SavedModel sul sito web di TensorFlow.

Per creare lo script di addestramento, esegui il seguente codice nel tuo 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"))

Una volta creato, lo script viene visualizzato nella cartella principale del notebook:

Visualizza lo script di addestramento.

Definisci gli argomenti per lo script di addestramento

Passa i seguenti argomenti della riga di comando allo script di addestramento:

  • label_column: identifica la colonna nei dati che contiene ciò che vuoi predire. In questo caso, la colonna è species. Lo hai definito in una variabile denominata LABEL_COLUMN durante l'elaborazione dei dati. Per ulteriori informazioni, consulta Scaricare, pre-elaborare e suddividere i dati.

  • epochs: il numero di epoche utilizzate per l'addestramento del modello. Un periodo è un'iterazione sui dati durante l'addestramento del modello. Questo tutorial utilizza 20 epoche.

  • batch_size: il numero di campioni elaborati prima degli aggiornamenti del modello. Questo tutorial utilizza una dimensione del batch pari a 10.

Per definire gli argomenti che vengono passati allo script, esegui il seguente codice:

JOB_NAME = "custom_job_unique"

EPOCHS = 20
BATCH_SIZE = 10

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