Cloud TPU(TF 2.x)での Transformer のトレーニング


Cloud TPU に慣れていない場合は、クイックスタートを参照して、TPU VM の作成方法を学習することを強くおすすめします。

このチュートリアルでは、Cloud TPU で Transformer モデルをトレーニングする方法について説明します。Transformer は、注意機構を使用してシーケンス間の問題を解決するニューラル ネットワーク アーキテクチャです。従来のニューラル seq2seq モデルとは異なり、Transformer は回帰結合を伴いません。注意機構は、トークン間の依存関係を 2 つのシーケンスで学習します。注意の重みはシーケンス内のすべてのトークンに適用されるため、Transformer モデルは長距離の依存関係を簡単にキャプチャできます。

Transformer の全体的な構造は、標準のエンコーダ-デコーダ パターンに従います。エンコーダは、自己注意を使用して入力シーケンスの表現を計算します。デコーダは、出力シーケンスを一度に 1 つずつ生成します。その際、エンコーダ出力と以前のデコーダ出力トークンを入力として使用します。

このモデルでは、入力トークンと出力トークンにも埋め込みを適用し、定位置エンコードを追加します。位置エンコードは、各トークンの位置に関する情報を追加します。

目標

  • データセットとモデルの出力を格納する Cloud Storage バケットを作成します。
  • モデルのトレーニングに使用するデータセットをダウンロードして前処理します。
  • トレーニング ジョブを実行します。
  • 出力結果を確認します。

費用

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

  • Compute Engine
  • Cloud TPU

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

始める前に

このチュートリアルを開始する前に、Google Cloud プロジェクトが正しく設定されていることを確認します。

  1. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.
  2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  3. Make sure that billing is enabled for your Google Cloud project.

  4. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  5. Make sure that billing is enabled for your Google Cloud project.

  6. このチュートリアルでは、Google Cloud の課金対象となるコンポーネントを使用します。費用を見積もるには、Cloud TPU の料金ページを確認してください。不要な課金を回避するために、このチュートリアルを完了したら、作成したリソースを必ずクリーンアップしてください。

単一の Cloud TPU デバイスでトレーニングする

このセクションでは、単一デバイスのトレーニング用に、Cloud Storage バケットと TPU VM を設定する方法について説明します。

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

    Cloud Shell を開く

  2. プロジェクト ID 用の環境変数を設定します。

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

    gcloud config set project ${PROJECT_ID}
    

    このコマンドを新しい Cloud Shell VM で初めて実行すると、Authorize Cloud Shell ページが表示されます。ページの下部にある [Authorize] をクリックして、gcloud に認証情報を使用した Google Cloud API の呼び出しを許可します。

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

    サービス アカウントにより、Cloud TPU サービスが他の Google Cloud サービスにアクセスできるようになります。

    $ 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-central2 gs://bucket-name
    

    この Cloud Storage バケットには、モデルのトレーニングに使用するデータとトレーニング結果が格納されます。このチュートリアルで使用する TPU を設定する gcloud コマンドは、前の手順で設定した Cloud TPU サービス アカウントのデフォルトの権限を設定します。権限の詳細な設定が必要な場合は、アクセスレベル権限をご覧ください。

単一の Cloud TPU でモデルをトレーニングする

  1. gcloud コマンドを使用して Cloud TPU VM を起動します。

      $ gcloud compute tpus tpu-vm create transformer-tutorial \
         --zone=us-central2-b \
         --accelerator-type=v4-8 \
         --version=tpu-vm-tf-2.17.0-pjrt
    

    コマンドフラグの説明

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

    gcloud compute tpus tpu-vm ssh transformer-tutorial --zone=us-central2-b
    
  3. 環境変数をエクスポートします。

      (vm)$ export STORAGE_BUCKET=gs://bucket-name
      (vm)$ export SENTENCEPIECE_MODEL=sentencepiece
      (vm)$ export SENTENCEPIECE_MODEL_PATH=${STORAGE_BUCKET}/${SENTENCEPIECE_MODEL}.model
      (vm)$ export TFDS_DIR=${STORAGE_BUCKET}/tfds
      (vm)$ export PARAM_SET=big
      (vm)$ export TPU_NAME=local
      (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/transformer/model_${PARAM_SET}
      (vm)$  export PYTHONPATH="/usr/share/tpu/models:$PYTHONPATH"
    
  4. TPU を作成するときに、--version パラメータを -pjrt で終わるバージョンに設定した場合は、次の環境変数を設定して PJRT ランタイムを有効にします。

      (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
      (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so
    
  5. TensorFlow の要件をインストールします。

    (vm)$ pip3 install -r /usr/share/tpu/models/official/requirements.txt
    
  6. データセットをダウンロードして前処理する

      (vm)$ python3 -c "import tensorflow_datasets as tfds; tfds.load('wmt14_translate/de-en', split='train+validation', shuffle_files=True, download=True)"
      (vm)$ python3 /usr/share/tpu/models/official/nlp/data/train_sentencepiece.py --output_model_path=${SENTENCEPIECE_MODEL}
    
  7. データセットを Cloud Storage バケットにコピーします

      (vm)$ gsutil -m cp ${SENTENCEPIECE_MODEL}.model ${STORAGE_BUCKET}
      (vm)$ gsutil -m cp ${SENTENCEPIECE_MODEL}.vocab ${STORAGE_BUCKET}
      (vm)$ gsutil -m cp -r tensorflow_datasets/wmt14_translate ${TFDS_DIR}/wmt14_translate
    
  8. トレーニング ディレクトリに移動します

    (vm)$ cd /usr/share/tpu/models/
    
  9. トレーニング スクリプトを実行します

    (vm)$ python3 official/nlp/train.py \
      --tpu=${TPU_NAME} \
      --experiment=wmt_transformer/large \
      --mode=train_and_eval \
      --model_dir=${MODEL_DIR} \
      --params_override="runtime.distribution_strategy=tpu, task.train_data.tfds_data_dir=${TFDS_DIR}, task.validation_data.tfds_data_dir=${TFDS_DIR}, task.sentencepiece_model_path=${SENTENCEPIECE_MODEL_PATH}, trainer.train_steps=10000, trainer.validation_interval=10000"
    

    コマンドフラグの説明

    tpu
    Cloud TPU の名前。環境変数(TPU_NAME)を指定することで設定されます。
    experiment
    トレーニングするモデル。
    mode
    スクリプトを実行するモード。
    model_dir
    トレーニング中にチェックポイントとサマリーが保存される Cloud Storage バケット。既存のフォルダを使用して、同じサイズの TPU と TensorFlow バージョンで以前に作成されたチェックポイントを読み込むことができます。
    params_override
    モデル パラメータを設定します。

    デフォルトでは、モデルは 10,000 ステップごとに評価されます。次のパラメータを設定して、トレーニング ステップの数を増やしたり、評価を実行する頻度を指定したりできます。

    • train.train_steps: 実行するトレーニング ステップの合計数です。
    • trainer.validation_interval: 評価間に実行するトレーニング ステップの数です。

    トレーニングと評価は、v4-8 Cloud TPU 上で約 20 分かかります。トレーニングと評価が完了すると、次のようなメッセージが表示されます。

    I0208 20:57:19.309512 140039467895872 controller.py:310]  eval | step:  10000 | eval time:   69.2 sec | output: 
    {'bleu_score': 19.204771518707275,
     'sacrebleu_score': 18.307039308307356,
     'validation_loss': 2.0654342}
     eval | step:  10000 | eval time:   69.2 sec | output: 
    {'bleu_score': 19.204771518707275,
     'sacrebleu_score': 18.307039308307356,
     'validation_loss': 2.0654342}
    

    これで単一デバイスのトレーニングが完了しました。 単一デバイスの TPU リソースを削除するには、次の手順を使用します。

  10. Compute Engine インスタンスから接続を切断します。

    (vm)$ exit
    

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

  11. TPU リソースを削除します。

      $ gcloud compute tpus tpu-vm delete transformer-tutorial \
        --zone=us-central2-b
      

    コマンドフラグの説明

    zone
    Cloud TPU が存在するゾーン

この時点で、このチュートリアルを終了してクリーンアップすることも、Cloud TPU Pod でのモデルの実行を続行して調べることもできます。

Cloud TPU Pod でモデルのスケーリングする

Cloud TPU Pod でモデルをトレーニングするには、トレーニング スクリプトに変更を加える必要がある場合があります。詳細については、TPU Pod でのトレーニングをご覧ください。

TPU Pod のトレーニング

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

    Cloud Shell を開く

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

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

    gcloud config set project ${PROJECT_ID}
    

    このコマンドを新しい Cloud Shell VM で初めて実行すると、Authorize Cloud Shell ページが表示されます。ページの下部にある [Authorize] をクリックして、gcloud に認証情報を使用した API の呼び出しを許可します。

  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 バケットには、モデルのトレーニングに使用するデータとトレーニング結果が格納されます。このチュートリアルで使用する gcloud コマンドは、前の手順で設定した Cloud TPU サービス アカウントのデフォルトの権限を設定します。権限の詳細な設定が必要な場合は、アクセスレベル権限をご覧ください。

TPU VM リソースを起動する

  1. gcloud コマンドを使用して TPU VM Pod を起動します。このチュートリアルでは、v4-32 Pod を指定します。他の Pod オプションについては、TPU タイプ利用可能な TPU タイプのページをご覧ください。

    $ gcloud compute tpus tpu-vm create transformer-tutorial \
      --zone=us-central2-b \
      --accelerator-type=v3-32 \
      --version=tpu-vm-tf-2.17.0-pod-pjrt
    

    コマンドフラグの説明

    zone
    Cloud TPU を作成するゾーン
    accelerator-type
    アクセラレータ タイプでは、作成する Cloud TPU のバージョンとサイズを指定します。TPU のバージョンごとにサポートされているアクセラレータ タイプの詳細については、TPU のバージョンをご覧ください。
    version
    Cloud TPU ソフトウェアのバージョン
  2. 次の ssh コマンドを実行して TPU VM に接続します。VM にログインすると、シェル プロンプトが username@projectname から username@vm-name に変わります。

    gcloud compute tpus tpu-vm ssh transformer-tutorial --zone=us-central2-b
    
  3. TensorFlow の要件をインストールします。

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

Pod トレーニングを設定して起動する

  1. Cloud TPU の設定変数をエクスポートします。

    (vm)$ export PYTHONPATH="/usr/share/tpu/models:$PYTHONPATH"
    (vm)$ export STORAGE_BUCKET=gs://bucket-name
    (vm)$ export SENTENCEPIECE_MODEL=sentencepiece
    (vm)$ export SENTENCEPIECE_MODEL_PATH=${STORAGE_BUCKET}/${SENTENCEPIECE_MODEL}.model
    (vm)$ export TFDS_DIR=${STORAGE_BUCKET}/tfds
    (vm)$ export TPU_NAME=transformer-tutorial
    (vm)$ export PARAM_SET=big
    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/transformer/model_${PARAM_SET} 
    (vm)$ export TPU_LOAD_LIBRARY=0
    
  2. データセットをダウンロードします

    (vm)$ python3 -c "import tensorflow_datasets as tfds; tfds.load('wmt14_translate/de-en', split='train+validation', shuffle_files=True, download=True)"
    (vm)$ python3 /usr/share/tpu/models/official/nlp/data/train_sentencepiece.py --output_model_path=${SENTENCEPIECE_MODEL}

  3. Copy the dataset to Cloud Storage bucket

    (vm)$ gsutil -m cp ${SENTENCEPIECE_MODEL}.model ${STORAGE_BUCKET}
    (vm)$ gsutil -m cp ${SENTENCEPIECE_MODEL}.vocab ${STORAGE_BUCKET}
    (vm)$ gsutil -m cp -r tensorflow_datasets/wmt14_translate ${TFDS_DIR}/wmt14_translate
    

  4. トレーニング ディレクトリに移動します。

    (vm)$ cd /usr/share/tpu/models/
  5. トレーニング スクリプトを実行します。

    (vm)$ python3 official/nlp/train.py 
    --tpu=${TPU_NAME}
    --experiment=wmt_transformer/large
    --mode=train_and_eval
    --model_dir=${MODEL_DIR}
    --params_override="runtime.distribution_strategy=tpu, task.train_data.tfds_data_dir=${TFDS_DIR}, task.validation_data.tfds_data_dir=${TFDS_DIR}, task.sentencepiece_model_path=${SENTENCEPIECE_MODEL_PATH}, trainer.train_steps=10000, trainer.validation_interval=10000"

    Command flag descriptions

    tpu
    The name of the Cloud TPU. This is set by specifying the environment variable (TPU_NAME).
    experiment
    The model to train.
    mode
    The mode in which to run the script.
    model_dir
    The Cloud Storage bucket where checkpoints and summaries are stored during training. You can use an existing folder to load previously generated checkpoints created on a TPU of the same size and TensorFlow version.
    params_override
    Set model parameters.

By default, the model will evaluate after every 10000 steps. In order to train to convergence, change train_steps to 200000. You can increase the number of training steps or specify how often to run evaluations by setting these parameters:

  • trainer.train_steps: Sets the total number of training steps to run.
  • trainer.validation_interval: Sets the number of training steps to run between evaluations.

Training and evaluation takes approximately 14 minutes on a v4-32 Cloud TPU. When the training and evaluation complete, messages similar to the following appear:

I0209 22:19:49.143219 139751309618240 controller.py:310]  eval | step:  10000 | eval time:   73.6 sec | output: 
    {'bleu_score': 19.401752948760986,
     'sacrebleu_score': 18.442741330886378,
     'validation_loss': 2.0558002}
 eval | step:  10000 | eval time:   73.6 sec | output: 
    {'bleu_score': 19.401752948760986,
     'sacrebleu_score': 18.442741330886378,
     'validation_loss': 2.0558002}

このトレーニングスクリプトは、20,000 ステップのトレーニングを行い、2,000 ステップごとに評価を実行します。このトレーニングと評価は、v3-32 Cloud TPU Pod 上で約 8 分かかります。トレーニングと評価が完了すると、次のようなメッセージが表示されます。

INFO:tensorflow:Writing to file /tmp/tmpdmlanxcf
I0218 21:09:19.100718 140509661046592 translate.py:184] Writing to file /tmp/tmpdmlanxcf
I0218 21:09:28.043537 140509661046592 transformer_main.py:118] Bleu score (uncased): 1.799112930893898
I0218 21:09:28.043911 140509661046592 transformer_main.py:119] Bleu score (cased): 1.730366237461567

収束するようにトレーニングするには、train_steps を 200,000 に変更します。次のパラメータを設定して、トレーニング ステップの数を増やしたり、評価を実行する頻度を指定したりできます。

  • --train_steps: 実行するトレーニング ステップの合計数を設定します。
  • --steps_between_evals: 評価間に実行するトレーニング ステップの数です。

トレーニングと評価が完了すると、次のようなメッセージが表示されます。

0509 00:27:59.984464 140553148962624 translate.py:184] Writing to file /tmp/tmp_rk3m8jp
I0509 00:28:11.189308 140553148962624 transformer_main.py:119] Bleu score (uncased): 1.3239131309092045
I0509 00:28:11.189623 140553148962624 transformer_main.py:120] Bleu score (cased): 1.2855342589318752

クリーンアップ

このチュートリアルで使用したリソースについて、Google Cloud アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。

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

    (vm)$ exit
    

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

  2. Cloud TPU と Compute Engine リソースを削除します。

    $ gcloud compute tpus tpu-vm delete transformer-tutorial \
      --zone=us-central2-b
    
  3. 次に示すように gsutil を実行します。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 推論の概要をご覧ください。