Vertex AI: Vertex AI Training で自動パッケージ化を使用して Hugging Face で BERT を微調整

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 ワークフローをサポートするさまざまなプロダクトが含まれています。このラボでは、主にトレーニングとワークベンチを使用します。

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 で置き換えてください。

ターミナルで gcloud config list --format 'value(core.project)' を実行すると、プロジェクト ID を取得できます。

PROJECT_ID='your-cloud-project'

次に、バケットを作成します。既存のバケットがある場合は、そちらを使用してください。

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

Vertex AI Training でカスタム トレーニング ジョブを実行する際は、GPU を使用します。ただし、Workbench インスタンスに GPU を指定していないため、ローカルテストには 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

要件ファイルを開き、次の内容を貼り付けます。

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 引数を見てみましょう。これは、カスタムジョブで使用するワーカープール構成を定義します。複数のワーカープールを使用するカスタムジョブを作成するには、複数のワーカープールの仕様を指定します。この例では、トレーニング コードが分散トレーニング用に構成されていないため、ワーカープールを 1 つだけ指定します。

この仕様の主なフィールドは次のとおりです。

• machine-type（必須）: マシンのタイプ。サポートされているタイプについては、こちらをご覧ください。
• replica-count: このワーカープールに使用するワーカー レプリカの数。デフォルト値は 1 です。
• accelerator-type: GPU のタイプ。サポートされているタイプについては、こちらをご覧ください。この例では、NVIDIA Tesla V100 GPU を 1 つ指定しました。
• 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 のナビゲーション メニューで [ストレージ] を移動してバケットを選択し、[削除] をクリックします。