Cloud TPU による BERT の調整: 文と文のペア分類タスク（TF 2.x）

このチュートリアルでは、Cloud TPU で BERT モデルのトレーニングを行う方法について説明します。

BERT は事前トレーニングの言語表現モデルです。事前トレーニングとは、最初にウィキペディアなどの大規模なテキストソースを使用して BERT のトレーニングを行うことを意味します。その後、このトレーニング結果を質問応答や感情分析などの他の自然言語処理（NLP）タスクに適用できます。BERT と Cloud TPU を使用すると、さまざまな NLP モデルを 30 分ほどでトレーニングできます。

BERT の詳細については、以下のリソースをご覧ください。

目標

  • データセットとモデルの出力を格納する Cloud Storage バケットを作成します。
  • トレーニング ジョブを実行します。
  • 出力結果を確認します。

費用

このドキュメントでは、Google Cloud の次の課金対象のコンポーネントを使用します。

  • Compute Engine
  • Cloud TPU
  • Cloud Storage

料金計算ツールを使うと、予想使用量に基づいて費用の見積もりを生成できます。 新しい Google Cloud ユーザーは無料トライアルをご利用いただける場合があります。

始める前に

このセクションでは、Cloud Storage バケットと Compute Engine VM の設定について説明します。

  1. Cloud Shell ウィンドウを開きます。

    Cloud Shell を開く

  2. プロジェクト ID の変数を作成します。

    export PROJECT_ID=project-id

  3. Cloud TPU を作成するプロジェクトを使用するように Google Cloud CLI を構成します。

    gcloud config set project ${PROJECT_ID}

  4. Cloud TPU プロジェクトのサービス アカウントを作成します。

    gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID

    このコマンドでは、Cloud TPU サービス アカウントを次の形式で返します。

    service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com

  5. 次のコマンドを使用して Cloud Storage バケットを作成します。

    gsutil mb -p ${PROJECT_ID} -c standard -l us-central1 gs://bucket-name

    この Cloud Storage バケットには、モデルのトレーニングに使用するデータとトレーニング結果が格納されます。TPU（TPU ノード アーキテクチャの場合は gcloud compute tpus execution-groups create、TPU VM アーキテクチャの場合は gcloud compute tpus tpu-vm create）の作成に使用するコマンドにより、前のステップで設定した Cloud TPU サービス アカウント用のデフォルトの権限が設定されます。権限の詳細な設定が必要な場合は、アクセスレベル権限をご覧ください。

    バケットのロケーションは、Compute Engine（VM）および Cloud TPU ノードと同じリージョンにする必要があります。

  6. gcloud コマンドを使用して Compute Engine VM と Cloud TPU を起動します。使用するコマンドは、TPU VM を使用するか TPU ノードを使用するかによって異なります。2 つの VM アーキテクチャの詳細については、システム アーキテクチャをご覧ください。gcloud コマンドの詳細については、gcloud リファレンスをご覧ください。

    TPU VM

    $ gcloud compute tpus tpu-vm create bert-tutorial \
--zone=us-central1-b \
--accelerator-type=v3-8 \
--version=tpu-vm-tf-2.16.1-se

    コマンドフラグの説明

    zone
    Cloud TPU を作成するゾーン。
    accelerator-type
    アクセラレータ タイプでは、作成する Cloud TPU のバージョンとサイズを指定します。TPU のバージョンごとにサポートされているアクセラレータ タイプの詳細については、TPU のバージョンをご覧ください。
    version
    Cloud TPU ソフトウェアのバージョン。

    TPU ノード

    $ gcloud compute tpus execution-groups create \
  --name=bert-tutorial \
  --zone=us-central1-b \
  --tf-version=2.12.0 \
  --machine-type=n1-standard-1 \
  --accelerator-type=v3-8

    コマンドフラグの説明

    name
    作成する Cloud TPU の名前。
    zone
    Cloud TPU を作成するゾーン。
    tf-version
    Tensorflow ctpu のバージョンが VM にインストールされます。
    machine-type
    作成する Compute Engine VM のマシンタイプ。
    accelerator type
    作成する Cloud TPU のタイプ。

  7. 自動的に Compute Engine インスタンスにログインしない場合は、次の ssh コマンドを実行してログインします。VM にログインすると、シェル プロンプトが username@projectname から username@vm-name に変わります。

    TPU VM

    gcloud compute tpus tpu-vm ssh bert-tutorial --zone=us-central1-b

    TPU ノード

    gcloud compute ssh bert-tutorial --zone=us-central1-b

    これらの手順を続行する場合は、VM セッション ウィンドウで、(vm)$ で始まる各コマンドを実行します。

  8. TPU 名の環境変数を作成します。

    TPU VM

    (vm)$ export TPU_NAME=local

    TPU ノード

    (vm)$ export TPU_NAME=bert-tutorial

データセットを準備する

  1. モデルとデータセットを保存するために必要なストレージ バケットを定義します。

    (vm)$ export STORAGE_BUCKET=gs://bucket-name

  2. トレーニング済みのチェックポイントとボキャブラリ ファイルをストレージ バケットにコピーします。

      (vm)$ curl https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/uncased_L-12_H-768_A-12.tar.gz -o uncased_L-12_H-768_A-12.tar.gz
  (vm)$ mkdir -p uncased_L-12_H-768_A-12
  (vm)$ tar -xvf uncased_L-12_H-768_A-12.tar.gz
  (vm)$ gsutil -m cp -R uncased_L-12_H-768_A-12 ${STORAGE_BUCKET}

モデルのトレーニング

  1. モデルのトレーニングと評価に必要なパラメータ値を定義します。

      (vm)$ export INIT_CHECKPOINT=${STORAGE_BUCKET}/uncased_L-12_H-768_A-12/bert_model.ckpt
  (vm)$ export TFDS_DIR=${STORAGE_BUCKET}/tfds
  (vm)$ export VOCAB_FILE=${STORAGE_BUCKET}/uncased_L-12_H-768_A-12/vocab.txt
  (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/bert-output
  (vm)$ export TASK=mnli

  2. TensorFlow の要件をインストールします。

    使用するコマンドは、TPU VM を使用するか TPU ノードを使用するかによって異なります。

    TPU VM

    (vm)$ pip3 install -r /usr/share/tpu/models/official/requirements.txt

    TPU ノード

    (vm)$ pip3 install --user -r /usr/share/models/official/requirements.txt
(vm)$ pip3 install tensorflow-datasets==4.6.0

  3. PYTHONPATH 環境変数を設定する

    TPU VM

    (vm)$ export PYTHONPATH=/usr/share/tpu/models

    TPU ノード

    (vm)$ export PYTHONPATH="${PYTHONPATH}:/usr/share/models"

  4. モデルが保存されているディレクトリに移動します。

    TPU VM

    (vm)$ cd /usr/share/tpu/models

    TPU ノード

    (vm)$ cd /usr/share/models

  5. トレーニング スクリプトを実行します。

    (vm)$ python3 official/nlp/train.py \
  --tpu=${TPU_NAME} \
  --experiment=bert/sentence_prediction_text \
  --mode=train_and_eval \
  --model_dir=${MODEL_DIR} \
  --config_file=official/nlp/configs/experiments/glue_mnli_text.yaml \
  --params_override="runtime.distribution_strategy=tpu, task.init_checkpoint=${INIT_CHECKPOINT}, task.train_data.tfds_data_dir=${TFDS_DIR}, task.train_data.vocab_file=${VOCAB_FILE}, task.validation_data.tfds_data_dir=${TFDS_DIR}, task.validation_data.vocab_file=${VOCAB_FILE}, trainer.train_steps=2000"

    コマンドフラグの説明

    tpu
    トレーニングに使用する Cloud TPU の名前です。
    mode
    train、eval、train_and_eval、predict のいずれか。
    model_dir
    モデルのトレーニングの際にチェックポイントとサマリーが保存される Cloud Storage パス。以前のチェックポイントが、同じサイズの Cloud TPU と TensorFlow バージョンを使用して作成されていれば、既存のフォルダを再利用して以前に生成されたチェックポイントを読み込み、追加のチェックポイントを保存できます。

    このスクリプトは、2,000 ステップのトレーニングを行った後、307 ステップの評価を実行します。v3-8 TPU では、トレーニング スクリプトが約 5 分後に完了し、次のような結果が表示されます。

    I0719 00:47:52.683979 140297079573568 controller.py:457] train | step:   2000 | steps/sec:   26.3 | output:
{'cls_accuracy': 0.7249375,
 'learning_rate': 1.4670059e-05,
 'training_loss': 0.6740678}
train | step:   2000 | steps/sec:   26.3 | output:
{'cls_accuracy': 0.7249375,
 'learning_rate': 1.4670059e-05,
 'training_loss': 0.6740678}
I0719 00:47:53.184051 140297079573568 controller.py:277]  eval | step:   2000 | running 307 steps of evaluation...
eval | step:   2000 | running 307 steps of evaluation...

クリーンアップ

  1. Compute Engine インスタンスとの接続を切断していない場合は切断します。

    (vm)$ exit

    プロンプトが username@projectname に変わります。これは、現在、Cloud Shell 内にいることを示しています。

  2. Cloud TPU と Compute Engine リソースを削除します。リソースの削除に使用するコマンドは、TPU VM または TPU ノードのどちらを使用するかによって異なります。詳細については、システム アーキテクチャをご覧ください。

    TPU VM

    $ gcloud compute tpus tpu-vm delete bert-tutorial \
--zone=us-central1-b

    TPU ノード

    $ gcloud compute tpus execution-groups delete bert-tutorial \
--zone=us-central1-b

  3. gcloud compute tpus execution-groups list を実行して、リソースが削除されたことを確認します。削除には数分かかることがあります。次のコマンドの出力には、このチュートリアルで作成したリソースを含めないでください。

    TPU VM

    $ gcloud compute tpus tpu-vm list --zone=us-central1-b

    TPU ノード

    $ gcloud compute tpus execution-groups list --zone=us-central1-b

  4. 次に示すように、gsutil を使用して Cloud Storage バケットを削除します。bucket-name を Cloud Storage バケットの名前に置き換えます。

    $ gsutil rm -r gs://bucket-name

次のステップ

TensorFlow Cloud TPU のチュートリアルでは通常、サンプル データセットを使用してモデルをトレーニングします。このトレーニングの結果は推論に使用できません。モデルを推論に使用するには、一般公開されているデータセットまたは独自のデータセットでデータをトレーニングします。Cloud TPU でトレーニングされた TensorFlow モデルは通常、データセットを TFRecord 形式にする必要があります。

データセット変換ツールのサンプルを使用して、画像分類データセットを TFRecord 形式に変換できます。画像分類モデルを使用しない場合は、自分でデータセットを TFRecord 形式に変換する必要があります。詳細については、TFRecord と tf.Example をご覧ください。

ハイパーパラメータ チューニング

データセットでモデルのパフォーマンスを向上させるには、モデルのハイパーパラメータを調整します。すべての TPU でサポートされているモデルに共通のハイパーパラメータに関する情報については、GitHub をご覧ください。モデルに固有のハイパーパラメータに関する情報については、各モデルのソースコードで確認できます。ハイパーパラメータ調整の詳細については、ハイパーパラメータ調整の概要とハイパーパラメータを調整するをご覧ください。

推論

モデルをトレーニングしたら、そのモデルを推論（予測）に使用できます。Cloud TPU 推論コンバータ ツールを使用して、Cloud TPU v5e での推論用の TensorFlow モデルを準備して最適化できます。Cloud TPU v5e での推論の詳細については、Cloud TPU v5e 推論の概要をご覧ください。