プロトタイプから本番環境へ: Vertex AI を使用したカスタムモデルのトレーニング

1. 概要

このラボでは、Vertex AI を使用して、カスタム トレーニング ジョブを実行します。

このラボは、動画シリーズ「プロトタイプから本番環境へ」の一部です。花のデータセットを使用して、画像分類モデルを構築します。詳細については、関連する動画でご確認ください。

学習内容

次の方法を学習します。

  • Vertex AI Workbench マネージド ノートブックを作成する
  • Vertex AI UI でカスタム トレーニング ジョブを構成して起動する
  • Vertex AI Python SDK でカスタム トレーニング ジョブを構成して起動する

このラボを Google Cloud で実行するための総費用は約 $1 です。

2. Vertex AI の概要

このラボでは、Google Cloud で利用できる最新の AI プロダクトを使用します。Vertex AI は Google Cloud 全体の ML サービスを統合してシームレスな開発エクスペリエンスを提供します。以前は、AutoML でトレーニングしたモデルやカスタムモデルには、個別のサービスを介してアクセスする必要がありました。Vertex AI は、これらの個別のサービスを他の新しいプロダクトとともに 1 つの API へと結合します。既存のプロジェクトを Vertex AI に移行することもできます。

Vertex AI には、エンドツーエンドの ML ワークフローをサポートするさまざまなプロダクトが含まれています。このラボでは、以下でハイライト表示されているプロダクト(TrainingWorkbench)を中心に学習します。

Vertex プロダクトの概要

3. 環境を設定する

この Codelab を実行するには、課金が有効になっている Google Cloud Platform プロジェクトが必要です。プロジェクトを作成するには、こちらの手順を行ってください。

ステップ 1: Compute Engine API を有効にする

まだ有効になっていない場合は、[Compute Engine] に移動して [有効にする] を選択します。

ステップ 2: Artifact Registry API を有効にする

まだ有効になっていない場合は、[Artifact Registry] に移動して [有効にする] を選択します。これは、カスタム トレーニング ジョブのコンテナを作成するために使用します。

ステップ 3: Vertex AI API を有効にする

Cloud コンソールの [Vertex AI] セクションに移動し、[Vertex AI API を有効にする] をクリックします。

Vertex AI ダッシュボード

ステップ 4: Vertex AI Workbench インスタンスを作成する

Cloud Console の [Vertex AI] セクションで [ワークベンチ] をクリックします。

Vertex AI メニュー

Notebooks API をまだ有効にしていない場合は、有効にします。

Notebook_api

有効にしたら、[マネージド ノートブック] をクリックします。

Notebooks_UI

[新しいノートブック] を選択します。

new_notebook

ノートブックの名前を指定して、[権限] で [Service account] を選択します。

create_notebook

[詳細設定] を選択します。

まだ有効になっていない場合は、[セキュリティ] で、[Enable terminal] を選択します。

enable_terminal

詳細設定のその他の設定はそのままで構いません。

[作成] をクリックします。インスタンスがプロビジョニングされるまでに数分かかります。

インスタンスが作成されたら、[JUPYTERLAB を開く] を選択します。

open_jupyterlab

4. トレーニング アプリケーション コードをコンテナ化する

このトレーニング ジョブを Vertex AI に送信しましょう。トレーニング アプリケーション コードを Docker コンテナに格納して、このコンテナを Google Artifact Registry に push します。この方法を使用すると、どのフレームワークで構築されたモデルでもトレーニングできます。

まず、Launcher メニューから、ノートブック インスタンスでターミナル ウィンドウを開きます。

ノートブックでターミナルを開く

ステップ 1: Cloud Storage バケットを作成する

このトレーニング ジョブでは、トレーニング済みの TensorFlow モデルを Cloud Storage バケットにエクスポートします。また、トレーニング用のデータも Cloud Storage バケットに保存します。

ターミナルで以下のように実行して、プロジェクトの環境変数を定義します。その際、your-cloud-project は実際のプロジェクト ID で置き換えてください。

PROJECT_ID='your-cloud-project'

次に、ターミナルで次のコマンドを実行して、プロジェクトに新しいバケットを作成します。

BUCKET="gs://${PROJECT_ID}-bucket"
gsutil mb -l us-central1 $BUCKET

ステップ 2: データを Cloud Storage バケットにコピーする

花のデータベースを Cloud Storage にコピーする必要があります。このデモでは、この Workbench インスタンスにデータセットをダウンロードして、それをバケットにコピーします。

データをダウンロードして解凍します。

wget https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz
tar xvzf flower_photos.tgz

解凍したデータをここで作成したバケットにコピーします。ディレクトリ全体をコピーしたいので -r を追加します。また、時間を短縮するため、-m を追加してマルチ処理でコピーを実行します。

gsutil -m cp -r flower_photos $BUCKET

ステップ 3: トレーニング コードを作成する

flowers という新しいディレクトリを作成し、そのディレクトリに移動します。

mkdir flowers
cd flowers

次のコマンドを実行して、トレーニング コードのディレクトリと、コードを追加する Python ファイルを作成します。

mkdir trainer
touch trainer/task.py

flowers/ ディレクトリに、次のものが作成されます。

+ trainer/
    + task.py

トレーニング アプリケーション コードの構造については、こちらのドキュメントをご覧ください。

いま作成した task.py ファイルを開いて、次のコードをコピーします。

{your-gcs-bucket} は、作成した Cloud Storage バケットの名前で置き換えます。

Cloud Storage FUSE ツールを使うと、ローカル ファイル システムのファイルにアクセスするのと同じように、Vertex AI Training のトレーニング ジョブから Cloud Storage のデータにアクセスできます。カスタム トレーニング ジョブを開始すると、ジョブは /gcs ディレクトリを確認します。ここにすべての Cloud Storage バケットがサブディレクトリとして含まれています。このため、トレーニング コードのデータパスは /gcs で始まっています。

import tensorflow as tf
import numpy as np
import os

## Replace {your-gcs-bucket} !!
BUCKET_ROOT='/gcs/{your-gcs-bucket}'

# Define variables
NUM_CLASSES = 5
EPOCHS=10
BATCH_SIZE = 32

IMG_HEIGHT = 180
IMG_WIDTH = 180

DATA_DIR = f'{BUCKET_ROOT}/flower_photos'

def create_datasets(data_dir, batch_size):
  '''Creates train and validation datasets.'''

  train_dataset = tf.keras.utils.image_dataset_from_directory(
    data_dir,
    validation_split=0.2,
    subset="training",
    seed=123,
    image_size=(IMG_HEIGHT, IMG_WIDTH),
    batch_size=batch_size)

  validation_dataset = tf.keras.utils.image_dataset_from_directory(
    data_dir,
    validation_split=0.2,
    subset="validation",
    seed=123,
    image_size=(IMG_HEIGHT, IMG_WIDTH),
    batch_size=batch_size)

  train_dataset = train_dataset.cache().shuffle(1000).prefetch(buffer_size=tf.data.AUTOTUNE)
  validation_dataset = validation_dataset.cache().prefetch(buffer_size=tf.data.AUTOTUNE)

  return train_dataset, validation_dataset

def create_model():
  '''Creates model.'''

  model = tf.keras.Sequential([
    tf.keras.layers.Resizing(IMG_HEIGHT, IMG_WIDTH),
    tf.keras.layers.Rescaling(1./255, input_shape=(IMG_HEIGHT, IMG_WIDTH, 3)),
    tf.keras.layers.Conv2D(16, 3, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Conv2D(32, 3, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(NUM_CLASSES, activation='softmax')
  ])
  return model

# CREATE DATASETS
train_dataset, validation_dataset = create_datasets(DATA_DIR, BATCH_SIZE)

# CREATE/COMPILE MODEL
model = create_model()
model.compile(optimizer=tf.keras.optimizers.Adam(),
              loss=tf.keras.losses.SparseCategoricalCrossentropy(),
              metrics=['accuracy'])

# TRAIN MODEL
history = model.fit(
  train_dataset,
  validation_data=validation_dataset,
  epochs=EPOCHS
)

# SAVE MODEL
model.save(f'{BUCKET_ROOT}/model_output')

ステップ 4: Dockerfile を作成する

コードをコンテナ化するには、Dockerfile ファイルを作成する必要があります。Dockerfile には、イメージの実行に必要なすべてのコマンドを含めます。必要なライブラリがすべてインストールされ、トレーニング コードのエントリ ポイントが設定されます。

ターミナルで、flowers ディレクトリのルートに空の Dockerfile を作成します。

touch Dockerfile

flowers/ ディレクトリに、次のものが作成されます。

+ Dockerfile
+ trainer/
    + task.py

Dockerfile を開き、次の内容をコピーします。

FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-8

WORKDIR /

# Copies the trainer code to the docker image.
COPY trainer /trainer

# Sets up the entry point to invoke the trainer.
ENTRYPOINT ["python", "-m", "trainer.task"]

このファイルに記述されているコマンドについて見てみましょう。

FROM コマンドにはベースイメージを指定します。これは、作成したイメージがビルドされる親イメージです。ここでは、ベースイメージとして Deep Learning Container TensorFlow Enterprise 2.8 GPU Docker イメージを使用します。Google Cloud の Deep Learning Containers には、ML およびデータ サイエンスの一般的なフレームワークが数多くプリインストールされています。

WORKDIR コマンドには、後続の命令を実行するイメージのディレクトリを指定します。

COPY コマンドは、トレーナー コードを Docker イメージにコピーします。この例では trainer ディレクトリに Python ファイルが 1 つしかありませんが、実際には他のファイルが存在している場合があります。たとえば、データ処理を扱う data.py や、モデルのコードを含む model.py などがあるかもしれません。より複雑なトレーニング コードについては、Python のドキュメントで Python プロジェクトのパッケージ化に関する説明をご覧ください。

別のライブラリを追加したい場合は、RUN コマンドを使用して pip install を実行します(例: RUN pip install -r requirements.txt)。ただし、ここでは何も追加する必要はありません。

最後の ENTRYPOINT コマンドは、トレーナーを呼び出すためのエントリ ポイントを設定します。トレーニング ジョブを開始したときに、これが実行されます。この例では、task.py ファイルを実行します。

Vertex AI Training 用の Dockerfile を記述する詳しい方法については、こちらをご覧ください。

ステップ 4: コンテナをビルドする

Workbench ノートブックのターミナルから次のコマンドを実行して、プロジェクトに env 変数を定義します。your-cloud-project はプロジェクトの ID で置き換えます。

PROJECT_ID='your-cloud-project'

Artifact Registry にリポジトリを作成します。

REPO_NAME='flower-app'

gcloud artifacts repositories create $REPO_NAME --repository-format=docker \
--location=us-central1 --description="Docker repository"

Google Artifact Registry 内のコンテナ イメージの URI を示す変数を定義します。

IMAGE_URI=us-central1-docker.pkg.dev/$PROJECT_ID/$REPO_NAME/flower_image:latest

Docker を構成します。

gcloud auth configure-docker \
    us-central1-docker.pkg.dev

続いて、flower ディレクトリのルートで次のように実行してコンテナをビルドします。

docker build ./ -t $IMAGE_URI

最後に、コンテナを Artifact Registry に push します。

docker push $IMAGE_URI

コンテナが Artifact Registry に push され、トレーニング ジョブを開始する準備ができました。

5. Vertex AI でカスタム トレーニング ジョブを実行する

このラボでは、Google Artifact Registry のカスタム コンテナによるカスタム トレーニングを使用します。トレーニング ジョブは、ビルド済みのコンテナを使って実行することもできます。

まず、Cloud コンソールの Vertex で [トレーニング] セクションに移動します。

トレーニング メニュー

ステップ 1: トレーニング ジョブを構成する

[作成] をクリックして、トレーニング ジョブのパラメータを入力します。

create_training

  • [Dataset] で [マネージド データセットなし] を選択します。
  • トレーニング方法として [カスタム トレーニング(上級者向け)] を選択し、[続行] をクリックします。
  • [新しいモデルのトレーニング] を選択して、モデルの名前に「flowers-model」と入力します(または任意のモデル名を入力します)。
  • [続行] をクリックします。

コンテナの設定ステップで、[カスタム コンテナ] を選択します。

カスタム コンテナ オプション

最初のボックス(コンテナ イメージ)で、前のセクションの IMAGE_URI 変数の値を入力します。これは us-central1-docker.pkg.dev/{PROJECT_ID}/flower-app/flower_image:latest のようになります(プロジェクト ID はご自分のプロジェクトの ID になります)。その他のフィールドは空白のままにして、[続行] をクリックします。

もう一度 [続行] をクリックして、ハイパーパラメータのステップをスキップします。

ステップ 2: コンピューティング クラスタを構成する

次のようにワーカープール 0 を構成します。

worker_pool_0

ステップ 6 はスキップします。予測コンテナはこのシリーズの次のラボで構成します。

[トレーニングを開始] をクリックして、トレーニング ジョブを開始します。コンソールの [トレーニング] セクションで、[TRAINING PIPELINES] タブの下に、新たに起動したジョブが表示されます。

トレーニング ジョブ

お疲れさまでした

Vertex AI を使って次のことを行う方法を学びました。

  • カスタム コンテナに用意されているトレーニング コード用のカスタム トレーニング ジョブを起動する。ここでは例として TensorFlow モデルを使用しましたが、カスタム コンテナまたはビルド済みコンテナを使って任意のフレームワークで構築されたモデルをトレーニングすることもできます。

Vertex のさまざまな部分の説明については、ドキュメントをご覧ください。

6. [省略可] Vertex AI Python SDK を使用する

前のセクションでは、UI からトレーニング ジョブを起動しました。このセクションでは、Vertex AI Python SDK を使用してトレーニング ジョブを送信します。

ノートブック インスタンスに戻り、Launcher から TensorFlow 2 ノートブックを作成します。

new_notebook

Vertex AI SDK をインポートします。

from google.cloud import aiplatform

CustomContainerTrainingJob を作成します。container_uri{PROJECT_ID} はプロジェクトの名前で、staging_bucket{BUCKET} は前に作成したバケットで置き換えます。

my_job = aiplatform.CustomContainerTrainingJob(display_name='flower-sdk-job',
                                               container_uri='us-central1-docker.pkg.dev/{PROJECT_ID}/flower-app/flower_image:latest',
                                               staging_bucket='gs://{BUCKET}')

ジョブを実行します。

my_job.run(replica_count=1,
           machine_type='n1-standard-8',
           accelerator_type='NVIDIA_TESLA_V100',
           accelerator_count=1)

デモの目的から、このジョブは前のセクションと比べるとサイズの大きいマシンで実行されるように構成されています。また、GPU を使用しています。machine-typeaccelerator_type、または accelerator_count を指定しないと、ジョブはデフォルトの n1-standard-4 で実行されます。

Cloud コンソールの [トレーニング] セクションにある [CUSTOM JOBS] タブに、トレーニング ジョブが表示されます。

7. クリーンアップ

Vertex AI Workbench マネージド ノートブックにはアイドル シャットダウン機能があります。インスタンスのシャットダウンを気にする必要はありません。インスタンスを手動でシャットダウンする場合は、コンソールで [Vertex AI] の [ワークベンチ] セクションにある [停止] ボタンをクリックします。ノートブックを完全に削除する場合は、[削除] ボタンをクリックします。

インスタンスの停止

ストレージ バケットを削除するには、Cloud コンソールのナビゲーション メニューで [ストレージ] に移動してバケットを選択し、[削除] をクリックします。

ストレージを削除