1. 概要
このラボでは、自動パッケージ化機能を使用して Vertex AI Training でカスタム トレーニング ジョブを実行する方法について学習します。Vertex AI のカスタム トレーニング ジョブはコンテナを使用します。独自の画像をビルドしない場合は、自動パッケージ化を使用できます。これにより、コードに基づいてカスタム Docker イメージがビルドされ、イメージが Container Registry に push され、イメージに基づいて CustomJob が開始されます。
学習内容
次の方法を学習します。
このラボを Google Cloud で実行するための総費用は約 $2 です。
2. ユースケースの概要
Hugging Face のライブラリを使用して、IMDB データセットで BERT モデルをファインチューニングします。このモデルは、映画のレビューがポジティブかネガティブかを予測します。データセットは Hugging Face データセット ライブラリから、Bert モデルは Hugging Face Transformers ライブラリからダウンロードされます。
3. Vertex AI の概要
このラボでは、Google Cloud で利用できる最新の AI プロダクトを使用します。Vertex AI は Google Cloud 全体の ML サービスを統合してシームレスな開発エクスペリエンスを提供します。以前は、AutoML でトレーニングしたモデルやカスタムモデルには、個別のサービスを介してアクセスする必要がありました。Vertex AI は、これらの個別のサービスを他の新しいプロダクトとともに 1 つの API へと結合します。既存のプロジェクトを Vertex AI に移行することもできます。ご意見やご質問がありましたら、サポートページからお寄せください。
Vertex AI には、エンドツーエンドの ML ワークフローをサポートするさまざまなプロダクトが含まれています。このラボでは、Training と 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 のデータセットと Transformers ライブラリを使用します。これらのライブラリは、ベースイメージとして選択したイメージに含まれていないため、要件として指定する必要があります。これを行うには、autopkg-codelab ディレクトリに requirements.txt ファイルを作成します。
autopkg-codelab ディレクトリにいることを確認して、ターミナルに次のように入力します。
touch requirements.txt
autopkg-codelab ディレクトリに、次のものが作成されます。
+ requirements.txt
+ trainer/
+ task.py
requirements ファイルを開き、次の内容を貼り付けます。
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 つ上の階層に移動します。
+ 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 ファイルで指定されているように、データセットと Transformers ライブラリを 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 引数を見てみましょう。これにより、カスタムジョブで使用するワーカープール構成が定義されます。分散トレーニング用の複数のワーカープールを使用するカスタムジョブを作成するために、複数のワーカープールの仕様を指定できます。この例では、トレーニング コードが分散トレーニング用に構成されていないため、単一のワーカープールのみを指定します。
この仕様の主なフィールドは次のとおりです。
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 のナビゲーション メニューで [ストレージ] を移動してバケットを選択し、[削除] をクリックします。
