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

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

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

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

費用

このチュートリアルでは、Google Cloud の課金対象となる以下のコンポーネントを使用します。

  • Compute Engine
  • Cloud TPU

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

始める前に

TPU Pod スライスのトレーニングを行う場合は、このドキュメントをお読みいただき、Pod スライスのトレーニングについて特に留意すべき事項をご確認ください。

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

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

    Cloud Shell を開く

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

    export PROJECT_ID=project-id

  3. Cloud TPU を作成するプロジェクトを使用するように gcloud コマンドライン ツールを構成します。

    gcloud config set project ${PROJECT_ID}

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

    gsutil mb -p ${PROJECT_ID} -c standard -l europe-west4 -b on gs://bucket-name

  5. ctpu up コマンドを使用して、Compute Engine VM を起動します。この例では、ゾーンを europe-west4-a に設定しますが、Compute Engine VM と Cloud TPU で使用する予定の任意のゾーンに設定できます。

    ctpu up --vm-only \
 --disk-size-gb=300 \
 --machine-type=n1-standard-8 \
 --zone=europe-west4-a \
 --tf-version=2.2 \
 --name=transformer-tutorial

  6. 指定した構成が表示されます。承認する場合は y、キャンセルする場合は n を入力してください。

  7. ctpu up コマンドの実行が終了したら、shell プロンプトが username@projectname から username@vm-name に変更されたことを確認します。変更されていれば、Compute Engine VM にログインしていることになります。

    gcloud compute ssh transformer-tutorial --zone=europe-west4-a

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

トレーニング データセットを生成する

Compute Engine VM の場合:

  1. 次の環境変数を作成します。

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

    (vm)$ export GCS_DATA_DIR=$STORAGE_BUCKET/data/transformer
(vm)$ export MODEL_DIR=$STORAGE_BUCKET/transformer/model_$PARAM_SET
(vm)$ export PYTHONPATH="$PYTHONPATH:/usr/share/models"
(vm)$ export PARAM_SET=big
(vm)$ export DATA_DIR=${HOME}/transformer/data
(vm)$ export VOCAB_FILE=${DATA_DIR}/vocab.ende.32768

  2. ディレクトリをトレーニング ディレクトリに変更します。

    (vm)$ cd /usr/share/models/official/nlp/transformer

  3. 次のデータセットの環境変数を設定します。

    


    (vm)$ export GCS_DATA_DIR=${STORAGE_BUCKET}/data/transformer
(vm)$ export MODEL_DIR=${STORAGE_BUCKET}/transformer/model_${PARAM_SET}

    4. 

  4. データセットをダウンロードして準備する
    


    (vm)$ python3 data_download.py --data_dir=${DATA_DIR}
(vm)$ gsutil cp -r ${DATA_DIR} ${GCS_DATA_DIR}

    


    data_download.py は、トレーニングと評価の WMT データセットをダウンロードして前処理します。データがダウンロードされて抽出されると、トレーニング データを使用してサブトークンのボキャブラリが生成されます。評価とトレーニングの文字列はトークン化されます。結果として出力されるデータはシャーディング、シャッフルされて TFRecord として保存されます。
    


    1.75 GB の圧縮データがダウンロードされます。圧縮ファイル、抽出ファイル、結合ファイルの合計ファイルサイズは 8.4 GB です。結果として生成される TFRecord ファイルとボキャブラリ ファイルは 722 MB です。スクリプトの実行には約 40 分かかります。その大部分がダウンロードに費やされ、前処理には 15 分ほどかかります。
    5. 




単一の Cloud TPU で英語 - ドイツ語の翻訳モデルをトレーニングする


Compute Engine VM で次のコマンドを実行します。



    

  1. 次のコマンドを実行して Cloud TPU を作成します。
    


    (vm)$ ctpu up --tpu-only \
  --tpu-size=v3-8  \
  --zone=europe-west4-a \
  --tf-version=2.2

    


    
  
    
    
      パラメータ
      説明
    
    
    
    
    
      tpu-size
    
    Cloud TPU のタイプ（v3-8 など）を指定します。
    
    
    
    
      zone
      Cloud TPU を作成するゾーン。これは、Compute Engine VM に使用したのと同じゾーンである必要があります。たとえば、europe-west4-a です。
      
    
    
      

    


    指定した構成が表示されます。承認する場合は y、キャンセルする場合は n を入力してください。
    


    Operation success; not ssh-ing to
Compute Engine VM
due to --tpu-only flag というメッセージが表示されます。SSH 鍵の伝播は完了しているため、このメッセージは無視できます。
    2. 

  2. Cloud TPU 名の変数を設定します。これは、--name パラメータで ctpu up に指定した名前か、デフォルトのユーザー名になります。
    


    (vm)$ export TPU_NAME=transformer-tutorial

    3. 

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


    (vm)$ python3 transformer_main.py \
    --tpu=${TPU_NAME} \
    --model_dir=${MODEL_DIR} \
    --data_dir=${GCS_DATA_DIR} \
    --vocab_file=${GCS_DATA_DIR}/vocab.ende.32768 \
    --bleu_source=${GCS_DATA_DIR}/newstest2014.en \
    --bleu_ref=${GCS_DATA_DIR}/newstest2014.de \
    --batch_size=6144 \
    --train_steps=2000 \
    --static_batch=true \
    --use_ctl=true \
    --param_set=big \
    --max_length=64 \
    --decode_batch_size=32 \
    --decode_max_length=97 \
    --padded_decode=true \
    --distribution_strategy=tpu

    4. 




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



    

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

  • --steps_between_evals: 評価間に実行するトレーニング ステップの数です。
    • 




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



INFO:tensorflow:Writing to file /tmp/tmpf2gn8vpa
I1125 21:22:30.734232 140328028010240 translate.py:182] Writing to file /tmp/tmpf2gn8vpa
I1125 21:22:42.785628 140328028010240 transformer_main.py:121] Bleu score (uncased): 0.01125154594774358
I1125 21:22:42.786558 140328028010240 transformer_main.py:122] Bleu score (cased): 0.01123994225054048




モデル評価中に BLEU スコアを計算する



モデルが評価されるとき、次のフラグを使用して BLEU を計算します。



    

  • --bleu_source: 翻訳するテキストを含むファイルへのパスです。
    • 

  • --bleu_ref: 参照訳を含むファイルへのパスです。
    • 




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



Cloud TPU Pod を使用したモデルのスケーリング



Cloud TPU Pod を使用してモデルをスケーリングすると、より迅速に結果を得ることができます。完全にサポートされているモデルは、次の Pod スライスを使用できます。



    

  • v2-32
    • 

  • v3-32
    • 




    

  1. Cloud Shell で、Cloud TPU の設定時に使用した tpu-only フラグと --zone フラグを指定して ctpu delete を実行します。これにより、Cloud TPU のみが削除されます。
    


     (vm)$ ctpu delete --tpu-only --zone=europe-west4-a
 
    2. 

  2. 使用する Pod スライスを指定するための tpu-sizeパラメータを使用して、ctpu up コマンドを実行します。たとえば、次のコマンドは v2-32 Pod スライスを使用します。
    


    (vm)$ ctpu up --tpu-only \
--tpu-size=v2-32  \
--zone=europe-west4-a \
--tf-version=2.2 \
--name=transformer-tutorial
    


    gcloud compute ssh transformer-tutorial --zone=europe-west4-a

    3. 

  3. TPU 名をエクスポートします。
    


    (vm)$ export TPU_NAME=transformer-tutorial

    4. 

  4. モデルディレクトリの変数をエクスポートします。
    


    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/transformer/model_${PARAM_SET}_pod

    5. 

  5. ディレクトリをトレーニング ディレクトリに変更します。
    


    (vm)$ cd /usr/share/models/official/nlp/transformer

    

    6. 

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


    (vm)$ python3 transformer_main.py \
     --tpu=${TPU_NAME} \
     --model_dir=${MODEL_DIR} \
     --data_dir=${GCS_DATA_DIR} \
     --vocab_file=${GCS_DATA_DIR}/vocab.ende.32768 \
     --bleu_source=${GCS_DATA_DIR}/newstest2014.en \
     --bleu_ref=${GCS_DATA_DIR}/newstest2014.de \
     --batch_size=24576 \
     --train_steps=2000 \
     --static_batch=true \
     --use_ctl=true \
     --param_set=big \
     --max_length=64 \
     --decode_batch_size=32 \
     --decode_max_length=97 \
     --padded_decode=true \
     --steps_between_evals=2000 \
     --distribution_strategy=tpu

    7. 




このトレーニングスクリプトは、2000 ステップのトレーニングを行い、2000 ステップごとに評価を実行します。このトレーニングと評価は、v2-32 Cloud TPU Pod 上で約 8 分かかります。



収束するようにトレーニングするには、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 Platform アカウントに課金されないようにする手順は次のとおりです。


1. Cloud Shell で、Cloud TPU の設定時に使用した --zone フラグを指定して ctpu delete を実行し、Compute Engine VM と Cloud TPU を削除します。
    $ ctpu delete --zone=europe-west4-a \
      --name=transformer-tutorial
    



    

  1. 次のコマンドを実行して、Compute Engine VM と Cloud TPU がシャットダウンされたことを確認します。
    


    $ ctpu status --zone=europe-west4-a

    


    削除には数分かかることがあります。次のようなレスポンスは、割り当てられたインスタンスがないことを示します。
    


    2018/04/28 16:16:23 WARNING: Setting zone to "europe-west4-a"
No instances currently exist.
        Compute Engine VM:     --
        Cloud TPU:             --

    2. 

  2. 次に示すように gsutil を実行します。bucket-name の部分は、このチュートリアルで作成した Cloud Storage バケット名に置き換えてください。
    


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

    3. 





    


    

      
次のステップ