トレーニング スクリプトを作成する

カスタムモデルを作成するには、カスタムモデルの作成とトレーニングを行う Python トレーニング スクリプトが必要です。Python トレーニング スクリプトでトレーニング ジョブを初期化してから、トレーニング ジョブの run メソッドを呼び出してスクリプトを実行します。

このトピックでは、トレーニング スクリプトを作成し、続いてトレーニング スクリプトのコマンド引数を指定します。

トレーニング スクリプトを作成する

このセクションでは、トレーニング スクリプトを作成します。このスクリプトは、お使いのノートブック環境内の task.py という名前の新しいファイルです。チュートリアルの後半では、このスクリプトを aiplatform.CustomTrainingJob コンストラクタに渡します。このスクリプトを実行すると、次の処理が行われます。

  • 作成した BigQuery データセットにデータを読み込みます。

  • TensorFlow Keras API を使用して、モデルの構築、コンパイル、トレーニングを行います。

  • Keras の Model.fit メソッドが呼び出されたときに使用するエポック数とバッチサイズを指定します。

  • AIP_MODEL_DIR 環境変数を使用して、モデル アーティファクトの保存場所を指定します。AIP_MODEL_DIR は Vertex AI によって設定され、モデル アーティファクトを保存するディレクトリの URI を含みます。詳細については、特別な Cloud Storage ディレクトリの環境変数をご覧ください。

  • TensorFlow の SavedModel をモデル ディレクトリにエクスポートします。詳細については、TensorFlow ウェブサイトの SavedModel 形式の使用をご覧ください。

トレーニング スクリプトを作成するには、ノートブックで次のコードを実行します。

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

スクリプトを作成すると、ノートブックのルートフォルダに表示されます。

トレーニング スクリプトを表示します。

トレーニング スクリプトの引数を定義する

トレーニング スクリプトに次のコマンドライン引数を渡します。

  • label_column - 予測対象を含むデータ内の列を特定します。この場合、その列は species です。これは、データを処理したときに LABEL_COLUMN という名前の変数で定義しました。詳細については、データのダウンロード、前処理、分割をご覧ください。

  • epochs - モデルのトレーニングに使用するエポックの数です。1 エポックとは、モデルをトレーニングする際にデータを 1 回反復処理することです。このチュートリアルでは、20 エポックを使用します。

  • batch_size - モデルを更新する前に処理されるサンプルの数です。このチュートリアルでは、バッチサイズ 10 を使用します。

スクリプトに渡される引数を定義するには、次のコードを実行します。

JOB_NAME = "custom_job_unique"

EPOCHS = 20
BATCH_SIZE = 10

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