Overview

Vertex AIを使って以下のことを行います。

モデルの学習 (Custom Training)
モデルの登録 (Vertex AI Model Registry)
モデルのDeploy (Deploy model to endpoint)
Deployしたモデルを使った推論 (Get prediction)

今回使う例

movielensのretrieve用のモデルを使うことにします。以下のページを参考に作成したものです。

import tensorflow_recommenders as tfrs

from typing import Dict, Text
import tempfile
import os
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds

# Env Var: https://cloud.google.com/vertex-ai/docs/training/code-requirements#environment-variables
MODEL_DIR = os.getenv("AIP_MODEL_DIR", tempfile.mkdtemp())
CHECKPOINT_DIR = os.path.join("AIP_CHECKPOINT_DIR", tempfile.mkdtemp())
TENSORBOARD_LOG_DIR = os.path.join("AIP_TENSORBOARD_LOG_DIR", tempfile.mkdtemp())

# Read data
ratings = tfds.load("movielens/100k-ratings", split="train")
# Features of all the available movies.
movies = tfds.load("movielens/100k-movies", split="train")

ratings = ratings.map(lambda x: {
    "movie_title": x["movie_title"],
    "user_id": x["user_id"]
})
movies = movies.map(lambda x: x["movie_title"]) # MapDataset で各ElementはTensor

user_ids_vocabulary = tf.keras.layers.StringLookup(mask_token=None)
user_ids_vocabulary.adapt(ratings.map(lambda x: x["user_id"]))

movie_titles_vocabulary = tf.keras.layers.StringLookup(mask_token=None)
movie_titles_vocabulary.adapt(movies)


class MovieLensModel(tfrs.Model):
  # We derive from a custom base class to help reduce boilerplate. Under the hood,
  # these are still plain Keras Models.

  def __init__(
      self,
      user_model: tf.keras.Model,
      movie_model: tf.keras.Model,
      task: tfrs.tasks.Retrieval):
    super().__init__()

    # Set up user and movie representations.
    self.user_model = user_model
    self.movie_model = movie_model

    # Set up a retrieval task.
    self.task = task

  def compute_loss(self, features: Dict[Text, tf.Tensor], training=False) -> tf.Tensor:
    # Define how the loss is computed.

    user_embeddings = self.user_model(features["user_id"])
    movie_embeddings = self.movie_model(features["movie_title"])

    return self.task(user_embeddings, movie_embeddings)

# Define user (user_id) and movie (movie_title) models.
user_model = tf.keras.Sequential([
    user_ids_vocabulary,
    tf.keras.layers.Embedding(user_ids_vocabulary.vocabulary_size(), 64)
])
movie_model = tf.keras.Sequential([
    movie_titles_vocabulary,
    tf.keras.layers.Embedding(movie_titles_vocabulary.vocabulary_size(), 64)
])

# Define your objectives.
task = tfrs.tasks.Retrieval(metrics=tfrs.metrics.FactorizedTopK(
    movies.batch(128).map(movie_model)
  )
)

# Create a retrieval model.
model = MovieLensModel(user_model, movie_model, task)
model.compile(optimizer=tf.keras.optimizers.Adagrad(0.5))

# Train for 3 epochs.
model.fit(ratings.batch(4096), epochs=3)


# !pip install -q scann
is_scann = False
try:
  index = tfrs.layers.factorized_top_k.ScaNN(model.user_model)
  index.index_from_dataset(
    tf.data.Dataset.zip((movies.batch(100), movies.batch(100).map(model.movie_model)))
  )
  is_scann = True
except:
  # Use brute-force search to set up retrieval using the trained representations.
  index = tfrs.layers.factorized_top_k.BruteForce(model.user_model)
  index.index_from_dataset(
    movies.batch(100).map(lambda title: (title, model.movie_model(title))))


# Get recommendations.
_, titles = index(np.array(["42"]))
print(f"Top 3 recommendations for user 42: {titles[0, :3]}")


index.save(MODEL_DIR, options=tf.saved_model.SaveOptions(namespace_whitelist=["Scann"]) if is_scann else None)
print(f"Model saved to {MODEL_DIR}")

ステップ

準備

GCPを使うので使う環境変数を設定しておきます。

export PROJECT=<your project>

適宜名前を変えていただければと思います。

export REGION=asia-northeast1
export REPOSITORY=ml-training
export IMAGE=movielens-retrieve
export IMAGE_TAG=0.0.1

モデルの学習

学習の仕方は複数あり、ローカルで学習しても、Cloud Run上で学習しても、Vertex AI のCustom Jobとして学習することも可能です。
上のScriptを実行するとモデルの学習が可能です。

最終的に、GCSなどにModelを保存することができれば大丈夫です。

今回は、2通りを実施しました。

モデル学習用のCustom Containerの作成 - 共通

Cloud Run、Vertex AI Custom Job両方で使うモデル学習用のカスタムコンテナを作成します。

今回はCloud buildを使用しました。（もちろん、LocalでBuildしたり、GitHub ActionsでBuild してPushしても大丈夫です。）

cloudbuild.yaml

steps:
  # build image for x86_64（amd64）platform
  - name: 'gcr.io/cloud-builders/docker'
    args:
      - buildx
      - build
      - --platform
      - linux/amd64
      - -t
      - ${_REGION}-docker.pkg.dev/${_PROJECT}/${_REPOSITORY}/${_IMAGE_NAME}:${_IMAGE_TAG}
      - .
    env:
      - 'DOCKER_CLI_EXPERIMENTAL=enabled'
  # push image to GAR
  - name: 'gcr.io/cloud-builders/docker'
    args:
      - push
      - ${_REGION}-docker.pkg.dev/${_PROJECT}/${_REPOSITORY}/${_IMAGE_NAME}:${_IMAGE_TAG}

images:
  - ${_REGION}-docker.pkg.dev/${_PROJECT}/${_REPOSITORY}/${_IMAGE_NAME}:${_IMAGE_TAG}

FROM python:3.12-slim

RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential \
    libhdf5-dev \
    pkg-config \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

# this is necessary ref: https://github.com/tensorflow/recommenders/issues/712
ENV TF_USE_LEGACY_KERAS=1

WORKDIR /app

COPY requirements.txt .

RUN pip install --no-cache-dir -r requirements.txt

COPY retrieve.py .

CMD ["python", "retrieve.py"]

gcloud builds submit \
    --config "cloudbuild.yaml" \
    --project "${PROJECT}" \
    --substitutions="_IMAGE_TAG=${IMAGE_TAG},_IMAGE_NAME=${IMAGE},_REPOSITORY=${REPOSITORY},_REGION=${REGION},_PROJECT=${PROJECT}" \
    --gcs-source-staging-dir="gs://${PROJECT}-cloudbuild/source"

モデル学習 - Cloud Run

Cloud Run JobでビルドしたDocker imageを実行します。

Cloud Run Jobの登録:

gcloud run jobs deploy ml-training-movielens-retrieve --memory 4Gi --cpu 2 --image "$REGION-docker.pkg.dev/$PROJECT/$REPOSITORY/$IMAGE:$IMAGE_TAG" --set-env-vars=AIP_MODEL_DIR=gs://${PROJECT}-ml-training/movielens/cloudrun/model-output --set-env-vars=TF_USE_LEGACY_KERAS=1 --max-retries 0 --region $REGION --project $PROJECT

AIP_MODEL_DIRという環境変数でアウトプットのGCSのPathを指定しています。

Cloud Run Jobの実行

gcloud run jobs execute ml-training-movielens-retrieve --region $REGION --project $PROJECT

結果の確認:

GCSの指定した場所にmodelのArtifactが格納されています。

gcloud storage ls "gs://${PROJECT}-ml-training/movielens/cloudrun/model-output/"
gs://PROJECT-ml-training/movielens/cloudrun/model-output/

gs://PROJECT-ml-training/movielens/cloudrun/model-output/:
gs://PROJECT-ml-training/movielens/cloudrun/model-output/
gs://PROJECT-ml-training/movielens/cloudrun/model-output/fingerprint.pb
gs://PROJECT-ml-training/movielens/cloudrun/model-output/keras_metadata.pb
gs://PROJECT-ml-training/movielens/cloudrun/model-output/saved_model.pb
gs://PROJECT-ml-training/movielens/cloudrun/model-output/assets/
gs://PROJECT-ml-training/movielens/cloudrun/model-output/model/
gs://PROJECT-ml-training/movielens/cloudrun/model-output/variables/

モデル学習 - VertexAI Custom Job

環境変数で、値を埋められるようにVertexAI のConfigファイルのTemplateを作成しました。 (直接Yamlを書いて実行しても問題ありません。)

vertexaiconfig.template.yaml

# https://cloud.google.com/vertex-ai/docs/reference/rest/v1/CustomJobSpec
workerPoolSpecs:
  machineSpec:
    machineType: n2-standard-2
  replicaCount: 1
  containerSpec:
    imageUri: $REGION-docker.pkg.dev/$PROJECT/$REPOSITORY/$IMAGE:$IMAGE_TAG
baseOutputDirectory:
  outputUriPrefix: gs://${PROJECT}-ml-training/movielens/vertexai/model-output/

個々では、baseOutputDirectory.outputUriPrefixにOutputのGCSのPathを指定します。

Env Varで実際のvertexaiconfig.yamlを生成します。

envsubst < tensorflow/examples/movielens/vertexaiconfig.template.yaml > tensorflow/examples/movielens/vertexaiconfig.yaml

Vertex AIのCustom Jobとして、実行します。

gcloud ai custom-jobs create --region=$REGION --display-name="movielens-retrieve" --config=tensorflow/examples/movielens/vertexaiconfig.yaml --project $PROJECT

同様にOutputにモデルのArtifactが格納されます。

モデルの登録

GCSに保存したモデルをVertex AI Model Registryに登録して、デプロイできるようにします。

gcloud ai models upload \
  --region=$REGION \
  --display-name=movielens-retrieve \
  --container-image-uri=asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.nightly:latest \
  --artifact-uri=gs://${PROJECT}-ml-training/movielens/vertexai/model-output/model/ \
  --project=$PROJECT

container-image-uri: モデルのServingで使われるImageです。今回はTensorflowのモデルなのでtensorflowのimageを使用しています。
display-name: 今回は movielens-retrieveとしました。
--artifact-uri: Cloud RunまたはVertex AI Custom Jobなどで学習したモデルのArtifactが保存されているDirectoryを指定します。

モデルのDeploy

モデルのデプロイをするためにEndpointを作成します。
Endpointは複数のモデルをServingしたり、特定のモデルだけをServingするように選択することができます。

curl -X POST \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
-H "Content-Type: application/json" \
-d '{"display_name": "movielens-retrieve", "dedicatedEndpointEnabled": true}' \
https://$REGION-aiplatform.googleapis.com/v1/projects/$PROJECT/locations/$REGION/endpoints

時間がかかる場合は、以下のgcloud ai endpoints createで作成したほうが良いかもしれません。

gcloud ai endpoints create \
  --region=$REGION \
  --display-name=movielens-retrieve --project $PROJECT

エンドポイントの確認

gcloud ai endpoints list --region=$REGION --project $PROJECT
ENDPOINT_ID          DISPLAY_NAME
xxxxxxxxxxxxxxxxxxx  movielens-retrieve

モデルをEndpointにDeployします。

ENDPOINT=$(gcloud ai endpoints list --region=$REGION --filter=display_name=movielens-retrieve --project $PROJECT --format="json(name)" | jq -r '.[0].name')
MODEL_ID=$(gcloud ai models list --filter=display_name=movielens-retrieve --region $REGION --project $PROJECT --format 'json(name)' | jq -r '.[0].name' | sed 's/.*\/\(\d*\)/\1/')

DeployするときにMachine Typeを選択することができます。実験の場合には小さめのインスタンスを選ぶとコストが抑えられます。

gcloud ai endpoints deploy-model $ENDPOINT \
  --region=$REGION \
  --model=$MODEL_ID \
  --display-name=movielens-retrieve \
  --machine-type=n2-standard-2 \
  --min-replica-count=1 \
  --max-replica-count=1 \
  --traffic-split=0=100 \
  --project $PROJECT

deployには5〜10分程度かかります。

推論 (Online Prediction)

推論するときに渡してあげる input_data_file.jsonを準備します。

今回のモデルではUserIdを与えるだけで結果を得ることができるので、user_idを配列で渡します。

{
    "instances": [
         "42"
    ]
}

EndpointIDを取得します。

ENDPOINT_ID=$(gcloud ai endpoints list --region=$REGION --filter=display_name=movielens-retrieve --project $PROJECT --format="json(name)" | jq -r '.[0].name' | sed 's/.*\/\(\d*\)/\1/')
INPUT_DATA_FILE=tensorflow/examples/movielens/input_data_file.json

Curlを使って、推論のEndpointに input data を渡して推論リクエストを送信します。

curl \
-X POST \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
-H "Content-Type: application/json" \
"https://$REGION-aiplatform.googleapis.com/v1/projects/$PROJECT/locations/$REGION/endpoints/$ENDPOINT_ID:predict" \
-d "@${INPUT_DATA_FILE}"

結果は以下のようになります:

{
  "predictions": [
    {
      "output_2": [
        "Rent-a-Kid (1995)",
        "Far From Home: The Adventures of Yellow Dog (1995)",
        "Just Cause (1995)",
        "Land Before Time III: The Time of the Great Giving (1995) (V)",
        "Nell (1994)",
        "Two if by Sea (1996)",
        "Jack (1996)",
        "Panther (1995)",
        "House Arrest (1996)",
        "Conan the Barbarian (1981)"
      ],
      "output_1": [
        3.94025946,
        3.47775483,
        3.4017539,
        3.32554197,
        2.95510435,
        2.63177681,
        2.61488819,
        2.61403036,
        2.58744907,
        2.54093599
      ]
    }
  ],
  "deployedModelId": "535000367843246080",
  "model": "projects/xxxx/locations/asia-northeast1/models/2548324905556901888",
  "modelDisplayName": "movielens-retrieve",
  "modelVersionId": "1"
}

dedicated endpointを使っていると以下のエラーメッセージが出ます。

{
  "error": {
    "code": 400,
    "message": "This endpoint is a dedicated endpoint via CloudESF and cannot be accessed through the Vertex AI API. Please access the endpoint using its dedicated dns name 'xxx.asia-northeast1-xxx.prediction.vertexai.goog'",
    "status": "FAILED_PRECONDITION"
  }
}

この場合は、指定されたDedicated DNS を使います。以下のようにendpointから取得することもできます。

DEDICATED_DNS=$(gcloud ai endpoints describe $ENDPOINT \
   --project=$PROJECT \
   --region=$REGION --format json | jq -r '.dedicatedEndpointDns')

curl \
    -X POST \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    -H "Content-Type: application/json" \
    "https://$DEDICATED_DNS/v1/projects/$PROJECT/locations/$REGION/endpoints/$ENDPOINT_ID:predict" \
    -d "@${INPUT_DATA_FILE}"

こちらはUI上からも実行することができます。

まとめ

Vertex AIを使ってレコメンデーションモデルを学習から推論するまでの過程を確認することができました。

Vertex AIでのその他のレコメンデーション

Vertex AI Agent Builderを使うと、自分でもモデルを学習する部分すらもManagedのモデルを使う事ができます。
モデルタイプやビジネス指標を選ぶだけで、レコメンデーションエンジンを作成しデプロイすることができます。

media recommendation を使った場合:

vertex ai search for retail を使った場合:

SearchをVertex AI Agent Builderを使って行った場合: