1. 概要
このラボでは、自動パッケージ化機能を使用して Vertex AI Training でカスタム トレーニング ジョブを実行する方法について学びます。Vertex AI のカスタム トレーニング ジョブはコンテナを使用します。独自のイメージをビルドしたくない場合は、auotpackaging を使用できます。これにより、コードに基づいてカスタム Docker イメージがビルドされ、Container Registry に push されて、イメージに基づく CustomJob が開始されます。
学習内容
次の方法を学習します。
このラボを Google Cloud で実行するための総費用は約 $2 です。
2. ユースケースの概要
Hugging Face のライブラリを使用して、IMDB データセットで BERT モデルをファインチューニングします。このモデルは、映画のレビューがポジティブかネガティブかを予測します。データセットは Hugging Face データセット ライブラリからダウンロードされ、BERT モデルは Hugging Face Transformer ライブラリからダウンロードされます。
3. Vertex AI の概要
このラボでは、Google Cloud で利用できる最新の AI プロダクトを使用します。Vertex AI は Google Cloud 全体の ML サービスを統合してシームレスな開発エクスペリエンスを提供します。以前は、AutoML でトレーニングしたモデルやカスタムモデルには、個別のサービスを介してアクセスする必要がありました。Vertex AI は、これらの個別のサービスを他の新しいプロダクトとともに 1 つの API へと結合します。既存のプロジェクトを Vertex AI に移行することもできます。ご意見やご質問がありましたら、サポートページからお寄せください。
Vertex AI には、エンドツーエンドの ML ワークフローをサポートするさまざまなプロダクトが含まれています。このラボでは、トレーニングと Workbench を中心に学習します。
4. 環境の設定
この Codelab を実行するには、課金が有効になっている Google Cloud Platform プロジェクトが必要です。プロジェクトを作成するには、こちらの手順を行ってください。
ステップ 1: Compute Engine API を有効にする
まだ有効になっていない場合は、[Compute Engine] に移動して [有効にする] を選択します。
ステップ 2: Vertex AI API を有効にする
Cloud コンソールの [Vertex AI] セクションに移動し、[Vertex AI API を有効にする] をクリックします。
ステップ 3: Container Registry API を有効にする
まだ有効になっていない場合は、[Container Registry] に移動して [有効にする] を選択します。これは、カスタム トレーニング ジョブのコンテナを作成するために使用します。
ステップ 4: Vertex AI Workbench インスタンスを作成する
Cloud Console の [Vertex AI] セクションで [ワークベンチ] をクリックします。
そこから、[マネージド ノートブック] をクリックします。
[新しいノートブック] を選択します。
ノートブックに名前を付けて、[詳細設定] をクリックします。
[詳細設定] で、アイドル状態でのシャットダウンを有効にして、シャットダウンまでの時間（分）を 60 に設定します。これにより、使用されていないノートブックが自動的にシャットダウンされるため、不要なコストが発生しません。
詳細設定のその他の設定はそのままで構いません。
次に、[作成] をクリックします。
インスタンスが作成されたら、[JUPYTERLAB を開く] を選択します。
新しいインスタンスを初めて使用するときは、認証を求められます。
5. トレーニング コードを作成する
まず、Launcher メニューから、ノートブック インスタンスでターミナル ウィンドウを開きます。
autopkg-codelab という新しいディレクトリを作成し、そのディスクに移動します。
mkdir autopkg-codelab
cd autopkg-codelab
ターミナルで、次のように実行してトレーニング コード用のディレクトリと、コードを追加する Python ファイルを作成します。
mkdir trainer
touch trainer/task.py
autopkg-codelab/ ディレクトリに、次のものが作成されます。
+ trainer/
+ task.py
いま作成した task.py ファイルを開いて、次のコードをコピーします。
import argparse
import tensorflow as tf
from datasets import load_dataset
from transformers import AutoTokenizer
from transformers import TFAutoModelForSequenceClassification
CHECKPOINT = "bert-base-cased"
def get_args():
'''Parses args.'''
parser = argparse.ArgumentParser()
parser.add_argument(
'--epochs',
required=False,
default=3,
type=int,
help='number of epochs')
parser.add_argument(
'--job_dir',
required=True,
type=str,
help='bucket to store saved model, include gs://')
args = parser.parse_args()
return args
def create_datasets():
'''Creates a tf.data.Dataset for train and evaluation.'''
raw_datasets = load_dataset('imdb')
tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
tokenized_datasets = raw_datasets.map((lambda examples: tokenize_function(examples, tokenizer)), batched=True)
# To speed up training, we use only a portion of the data.
# Use full_train_dataset and full_eval_dataset if you want to train on all the data.
small_train_dataset = tokenized_datasets['train'].shuffle(seed=42).select(range(1000))
small_eval_dataset = tokenized_datasets['test'].shuffle(seed=42).select(range(1000))
full_train_dataset = tokenized_datasets['train']
full_eval_dataset = tokenized_datasets['test']
tf_train_dataset = small_train_dataset.remove_columns(['text']).with_format("tensorflow")
tf_eval_dataset = small_eval_dataset.remove_columns(['text']).with_format("tensorflow")
train_features = {x: tf_train_dataset[x] for x in tokenizer.model_input_names}
train_tf_dataset = tf.data.Dataset.from_tensor_slices((train_features, tf_train_dataset["label"]))
train_tf_dataset = train_tf_dataset.shuffle(len(tf_train_dataset)).batch(8)
eval_features = {x: tf_eval_dataset[x] for x in tokenizer.model_input_names}
eval_tf_dataset = tf.data.Dataset.from_tensor_slices((eval_features, tf_eval_dataset["label"]))
eval_tf_dataset = eval_tf_dataset.batch(8)
return train_tf_dataset, eval_tf_dataset
def tokenize_function(examples, tokenizer):
'''Tokenizes text examples.'''
return tokenizer(examples['text'], padding='max_length', truncation=True)
def main():
args = get_args()
train_tf_dataset, eval_tf_dataset = create_datasets()
model = TFAutoModelForSequenceClassification.from_pretrained(CHECKPOINT, num_labels=2)
model.compile(
optimizer=tf.keras.optimizers.Adam(learning_rate=0.01),
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=tf.metrics.SparseCategoricalAccuracy(),
)
model.fit(train_tf_dataset, validation_data=eval_tf_dataset, epochs=args.epochs)
model.save(f'{args.job_dir}/model_output')
if __name__ == "__main__":
main()
コードについて、以下の点に注意してください。
CHECKPOINTは、ファインチューニングするモデルです。ここでは BERT を使用します。TFAutoModelForSequenceClassificationメソッドは、指定された言語モデル アーキテクチャと重みを TensorFlow で読み込み、ランダムに初期化された重みを持つ分類ヘッドを上部に追加します。この例では、バイナリ分類問題（陽性または陰性）があるため、この分類器にはnum_labels=2を指定します。
6. トレーニング コードをコンテナ化してローカルで実行する
gcloud ai custom-jobs local-run コマンドを使用すると、トレーニング コードに基づいて Docker コンテナ イメージをビルドし、ローカルマシンでコンテナとして実行できます。コンテナをローカルで実行すると、Vertex AI Training での実行と同様の方法でトレーニング コードを実行できます。Vertex AI でカスタム トレーニングを実行する前に、コードの問題をデバッグできます。
トレーニング ジョブでは、トレーニング済みモデルを Cloud Storage バケットにエクスポートします。ターミナルで以下のように実行して、プロジェクトの環境変数を定義します。その際、your-cloud-project は実際のプロジェクト ID で置き換えてください。
PROJECT_ID='your-cloud-project'
次に、バケットを作成します。既存のバケットがある場合は、代わりにそれを使用しても構いません。
BUCKET_NAME="gs://${PROJECT_ID}-bucket"
gsutil mb -l us-central1 $BUCKET_NAME
Vertex AI Training でカスタム トレーニング ジョブを実行する場合は、GPU を使用します。ただし、GPU を使用する Workbench インスタンスを指定しなかったため、ローカルテストには CPU ベースのイメージを使用します。この例では、Vertex AI Training のビルド済みコンテナを使用します。
次のコマンドを実行して、コンテナのベースとして使用する Docker イメージの URI を設定します。
BASE_CPU_IMAGE=us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-7:latest
次に、ローカルの実行コマンドでビルドされた Docker イメージの名前を設定します。
OUTPUT_IMAGE=$PROJECT_ID-local-package-cpu:latest
トレーニング コードでは Hugging Face データセットと Transformer ライブラリを使用しています。これらのライブラリは、ベースイメージとして選択したイメージに含まれていないため、要件として提供する必要があります。そのために、autopkg-codelab ディレクトリに requirements.txt ファイルを作成します。
autopkg-codelab ディレクトリにいることを確認し、ターミナルで次のように入力します。
touch requirements.txt
autopkg-codelab ディレクトリに、次のものが作成されます。
+ requirements.txt
+ trainer/
+ task.py
要件ファイルを開き、次のコードを貼り付けます。
datasets==1.18.2
transformers==4.16.2
最後に、gcloud ai custom-jobs local-run コマンドを実行して、Workbench マネージド インスタンスでトレーニングを開始します。
gcloud ai custom-jobs local-run \
--executor-image-uri=$BASE_CPU_IMAGE \
--python-module=trainer.task \
--output-image-uri=$OUTPUT_IMAGE \
-- \
--job_dir=$BUCKET_NAME
Docker イメージがビルドされていることがわかります。requirements.txt ファイルに追加した依存関係が pip でインストールされます。このコマンドを初めて実行する場合は、完了するまでに数分かかることがあります。イメージがビルドされると、task.py ファイルの実行が開始され、モデルのトレーニングが表示されます。次のように表示されます。
ローカルで GPU を使用していないため、モデルのトレーニングに時間がかかります。ジョブの完了を待つ代わりに、Ctrl+C キーを押してローカル トレーニングをキャンセルできます。
さらにテストを行いたい場合は、上でビルドしたイメージを再パッケージ化せずに直接実行することもできます。
gcloud beta ai custom-jobs local-run \
--executor-image-uri=$OUTPUT_IMAGE \
-- \
--job_dir=$BUCKET_NAME \
--epochs=1
7. カスタムジョブを作成する
ローカルモードをテストしたので、自動パッケージ化機能を使用して Vertex AI Training でカスタム トレーニング ジョブを起動します。1 つのコマンドで、次の処理を実行できます。
- コードに基づいてカスタム Docker イメージをビルドします。
- イメージを Container Registry に push します。
- 画像に応じて
CustomJobを開始します。
ターミナルに戻り、autopkg-codelab ディレクトリの 1 つ上のレベルに cd します。
+ autopkg-codelab
+ requirements.txt
+ trainer/
+ task.py
Vertex AI Training のビルド済み TensorFlow GPU イメージをカスタム トレーニング ジョブのベースイメージとして指定する。
BASE_GPU_IMAGE=us-docker.pkg.dev/vertex-ai/training/tf-gpu.2-7:latest
次に、gcloud ai custom-jobs create コマンドを実行します。まず、このコマンドはトレーニング コードに基づいてカスタム Docker イメージをビルドします。ベースイメージは、BASE_GPU_IMAGE に設定した Vertex AI Training のビルド済みコンテナです。その後、自動パッケージ化機能により、requirements.txt ファイルで指定されたデータセットと Transformer ライブラリが pip によってインストールされます。
gcloud ai custom-jobs create \
--region=us-central1 \
--display-name=fine_tune_bert \
--args=--job_dir=$BUCKET_NAME \
--worker-pool-spec=machine-type=n1-standard-4,replica-count=1,accelerator-type=NVIDIA_TESLA_V100,executor-image-uri=$BASE_GPU_IMAGE,local-package-path=autopkg-codelab,python-module=trainer.task
worker-pool-spec 引数を見てみましょう。カスタムジョブで使用されるワーカープール構成を定義します。分散トレーニング用の複数のワーカープールを使用するカスタムジョブを作成するために、複数のワーカープールの仕様を指定できます。この例では、トレーニング コードが分散トレーニング用に構成されていないため、ワーカープールを 1 つだけ指定します。
この仕様の主なフィールドは次のとおりです。
machine-type（必須）: マシンのタイプ。サポートされているタイプについては、こちらをクリックしてください。replica-count: このワーカープールに使用するワーカー レプリカの数。デフォルトの値は 1 です。accelerator-type: GPU のタイプ。サポートされているタイプについては、こちらをクリックしてください。この例では、1 つの NVIDIA Tesla V100 GPU を指定しました。accelerator-count: 使用するワーカープール内の各 VM の GPU の数。デフォルトの値は 1 です。executor-image-uri: 指定されたパッケージを実行するコンテナ イメージの URI。これはベースイメージに設定されます。local-package-path: トレーニング コードを含むフォルダのローカルパス。python-module: 指定されたパッケージ内で実行する Python モジュール名。
ローカル コマンドを実行したときと同様に、Docker イメージがビルドされ、トレーニング ジョブが開始されます。トレーニング ジョブの出力は表示されず、トレーニング ジョブが起動されたことを示す次のメッセージが表示されます。custom-jobs create コマンドの初回実行時は、イメージがビルドされて push されるまで数分かかる場合があります。
Cloud コンソールの Vertex AI Training セクションに戻り、[CUSTOM JOBS] にジョブが実行されていることを確認します。
ジョブが完了するまでに 20 分ほどかかります。
完了すると、バケットの model_output ディレクトリに次の保存済みモデル アーティファクトが表示されます。
お疲れさまでした
Vertex AI を使って次のことを行う方法を学びました。
- トレーニング コードをコンテナ化してローカルで実行する
- 自動パッケージ化を使用してトレーニング ジョブを Vertex AI Training に送信する
Vertex AI のさまざまな機能の詳細については、こちらのドキュメントをご覧ください。
8. クリーンアップ
ノートブックは、アイドル状態で 60 分が経過するとタイムアウトするように構成されています。このため、インスタンスのシャットダウンを心配する必要はありません。インスタンスを手動でシャットダウンする場合は、Console で [Vertex AI] の [ワークベンチ] セクションにある [停止] ボタンをクリックします。ノートブックを完全に削除する場合は、[削除] ボタンをクリックします。
ストレージ バケットを削除するには、Cloud Console のナビゲーション メニューで [ストレージ] を移動してバケットを選択し、[削除] をクリックします。
