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 プロジェクトが正しく設定されていることを確認します。

Google Cloud アカウントにログインします。Google Cloud を初めて使用する場合は、アカウントを作成して、実際のシナリオでの Google プロダクトのパフォーマンスを評価してください。新規のお客様には、ワークロードの実行、テスト、デプロイができる無料クレジット $300 分を差し上げます。

Google Cloud Console の [プロジェクトセレクタ] ページで、Google Cloud プロジェクトを選択または作成します。

プロジェクトセレクタに移動

Google Cloud プロジェクトで課金が有効になっていることを確認します。

Google Cloud Console の [プロジェクトセレクタ] ページで、Google Cloud プロジェクトを選択または作成します。

プロジェクトセレクタに移動

Google Cloud プロジェクトで課金が有効になっていることを確認します。

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

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

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

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

Cloud Shell を開く
プロジェクト ID 用の環境変数を設定します。
```
export PROJECT_ID=project-id
```
Cloud TPU を作成する Google Cloud プロジェクトを使用するように Google Cloud CLI を構成します。
```
gcloud config set project ${PROJECT_ID}
```
このコマンドを新しい Cloud Shell VM で初めて実行すると、Authorize Cloud Shell ページが表示されます。ページの下部にある [Authorize] をクリックして、gcloud に認証情報を使用した Google Cloud API の呼び出しを許可します。
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
```
次のコマンドを使用して Cloud Storage バケットを作成します。
```
$ gsutil mb -p ${PROJECT_ID} -c standard -l us-central2 gs://bucket-name
```
この Cloud Storage バケットには、モデルのトレーニングに使用するデータとトレーニング結果が格納されます。このチュートリアルで使用する TPU を設定する gcloud コマンドは、前の手順で設定した Cloud TPU サービスアカウントのデフォルトの権限を設定します。権限の詳細な設定が必要な場合は、アクセスレベル権限をご覧ください。

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

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.16.1-pjrt
```
コマンドフラグの説明

zone

Cloud TPU を作成するゾーン。

accelerator-type

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

version

Cloud TPU ソフトウェアのバージョン。

注: Google Cloud プロジェクトが複数ある場合は、--project フラグを使用して、Cloud TPU を作成する Google Cloud プロジェクト ID を指定する必要があります。
注: プロジェクトで初めて gcloud compute tpus を実行する場合、SSH 認証鍵の伝播や API の始動などの起動タスクの実行に約 5 分かかります。
次の ssh コマンドを実行して、Cloud TPU VM に接続します。

注: Cloud TPU VM に接続すると、シェルプロンプトが username@projectname から username@vm-name に変わります。
```
gcloud compute tpus tpu-vm ssh transformer-tutorial --zone=us-central2-b
```
重要なポイント: これ以降、接頭辞 (vm) $ は Compute Engine VM インスタンスでコマンドを実行する必要があることを意味します。

環境変数をエクスポートします。

  (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"

TPU を作成するときに、--version パラメータを -pjrt で終わるバージョンに設定した場合は、次の環境変数を設定して PJRT ランタイムを有効にします。
```
  (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
  (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so
```

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

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

データセットをダウンロードして前処理する

  (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}

データセットを 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

トレーニングディレクトリに移動します
```
(vm)$ cd /usr/share/tpu/models/
```
トレーニングスクリプトを実行します
```
(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 リソースを削除するには、次の手順を使用します。
Compute Engine インスタンスから接続を切断します。
```
(vm)$ exit
```
プロンプトが username@projectname に変わります。これは、現在、Cloud Shell 内にいることを示しています。
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 のトレーニング

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

Cloud Shell を開く
プロジェクト ID の変数を作成します。
```
export PROJECT_ID=project-id
```
Cloud TPU を作成するプロジェクトを使用するように Google Cloud CLI を構成します。
```
gcloud config set project ${PROJECT_ID}
```
このコマンドを新しい Cloud Shell VM で初めて実行すると、Authorize Cloud Shell ページが表示されます。ページの下部にある [Authorize] をクリックして、gcloud に認証情報を使用した API の呼び出しを許可します。
Cloud TPU プロジェクトのサービスアカウントを作成します。
```
gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID
```
このコマンドでは、Cloud TPU サービスアカウントを次の形式で返します。
```
service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
```
次のコマンドを使用して Cloud Storage バケットを作成するか、以前にプロジェクト用に作成したバケットを使用します。
```
gsutil mb -p ${PROJECT_ID} -c standard -l us-central1 gs://bucket-name
```
この Cloud Storage バケットには、モデルのトレーニングに使用するデータとトレーニング結果が格納されます。このチュートリアルで使用する gcloud コマンドは、前の手順で設定した Cloud TPU サービスアカウントのデフォルトの権限を設定します。権限の詳細な設定が必要な場合は、アクセスレベル権限をご覧ください。

重要: 新しい Cloud Storage バケットを作成した場合は、データセットをダウンロードして前処理する必要があります。詳細については、データセットの準備をご覧ください。

TPU VM リソースを起動する

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

注: 現在 TPU Pod を作成するのに十分な容量がない場合は、キューに格納されたリソースを使用してリクエストをキューに入れることができます。キューに格納されたリソースを使用すると、利用可能になったときに容量を受信できます。Cloud TPU リソースをキューに登録されたリソースとしてリクエストするには、代わりに gcloud alpha compute tpus queued-resources create コマンドを使用します。詳細については、キューに格納されたリソースの管理をご覧ください。
```
$ gcloud compute tpus tpu-vm create transformer-tutorial \
  --zone=us-central2-b \
  --accelerator-type=v3-32 \
  --version=tpu-vm-tf-2.16.1-pod-pjrt
```
コマンドフラグの説明

zone

Cloud TPU を作成するゾーン。

accelerator-type

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

version

Cloud TPU ソフトウェアのバージョン。
次の ssh コマンドを実行して TPU VM に接続します。VM にログインすると、シェルプロンプトが username@projectname から username@vm-name に変わります。
```
gcloud compute tpus tpu-vm ssh transformer-tutorial --zone=us-central2-b
```

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

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

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

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

データセットをダウンロードします

(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}

データセットを Cloud Storage バケットにコピーします

注: すでにデータセットをダウンロードして 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
```
トレーニングディレクトリに移動します。
```
(vm)$ cd /usr/share/tpu/models/
```

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

(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.

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

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

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

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 アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。

Compute Engine インスタンスとの接続を切断していない場合は切断します。
```
(vm)$ exit
```
プロンプトが username@projectname に変わります。これは、現在、Cloud Shell 内にいることを示しています。

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

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

次に示すように 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 推論の概要をご覧ください。

Cloud TPU（TF 2.x）での Transformer のトレーニング

目標

費用

始める前に

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

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

コマンドフラグの説明

コマンドフラグの説明

コマンドフラグの説明

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

TPU Pod のトレーニング

TPU VM リソースを起動する

コマンドフラグの説明

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

Command flag descriptions

クリーンアップ

次のステップ

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

推論

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