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 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 europe-west4 -b on gs://bucket-name
    

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

  6. 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.3.1 \
     --name=transformer-tutorial
    

    コマンドフラグの説明

    vm-only
    VM のみを作成します。デフォルトでは、ctpu up コマンドは VM と Cloud TPU を作成します。
    disk-size-gb
    VM のディスクのサイズ(GB)。
    machine_type
    ctpu up コマンドで作成される VM のマシンタイプ
    zone
    Cloud TPU を作成するゾーン
    tpu-size
    作成する Cloud TPU のタイプ
    tf-version
    Tensorflow ctpu のバージョンが VM にインストールされます。
    name
    作成する Cloud TPU の名前。

    ctpu ユーティリティの詳細については、CTPU リファレンスをご覧ください。

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

  8. 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. 次の環境変数を作成します。bucket-name は、使用するバケット名に置き換えてください。

    (vm)$ export STORAGE_BUCKET=gs://bucket-name
    
    (vm)$ export GCS_DATA_DIR=$STORAGE_BUCKET/data/transformer
    (vm)$ export PARAM_SET=big
    (vm)$ export MODEL_DIR=$STORAGE_BUCKET/transformer/model_$PARAM_SET
    (vm)$ export PYTHONPATH="$PYTHONPATH:/usr/share/models"
    (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. データセットをダウンロードして準備します。

    (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 分ほどかかります。

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

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

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

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

    コマンドフラグの説明

    tpu-only
    Cloud TPU のみを作成します。デフォルトでは、ctpu up コマンドは VM と Cloud TPU を作成します。
    tpu-size
    Cloud TPU のタイプ(v3-8 など)を指定します。
    zone
    Cloud TPU を作成するゾーン。これは、Compute Engine VM に使用したのと同じゾーンである必要があります。例: europe-west4-a
    tf-version
    Tensorflow ctpu のバージョンが VM にインストールされます。
    name
    作成する Cloud TPU の名前。

    ctpu ユーティリティの詳細については、CTPU リファレンスをご覧ください。

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

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

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

    (vm)$ export TPU_NAME=transformer-tutorial
    
  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

    コマンドフラグの説明

    tpu
    Cloud TPU の名前。環境変数(TPU_NAME)を指定することで設定されます。
    model_dir
    モデルのトレーニング中にチェックポイントとサマリーが保存されるディレクトリ。フォルダがない場合は、プログラムによって作成されます。Cloud TPU を使用する場合、model_dir を Cloud Storage パスにする必要があります(`gs://...`)。以前のチェックポイントが、同じサイズの Cloud TPU と TensorFlow バージョンを使用して作成されていれば、既存のフォルダを再利用して現在のチェックポイント データを読み込んで追加のチェックポイントを保存できます。
    data_dir
    トレーニング入力の Cloud Storage パス。この例では、fake_imagenet データセットに設定されています。
    vocab_file
    翻訳対象の用語が含まれるファイル。
    bleu_source
    翻訳対象のソース文が含まれるファイル。
    bleu_ref
    翻訳文のリファレンスが含まれるファイル。
    train_steps
    モデルをトレーニングするためのステップの数。1 つのステップで 1 つのデータバッチを処理します。これには、バックワード パスと逆伝播の両方が含まれます。
    batch_size
    トレーニングのバッチサイズ。
    static_batch
    データセット内のバッチに静的な形状があるかどうかを指定します。
    use_ctl
    スクリプトをカスタム トレーニング ループで実行するかどうかを指定します。
    param_set
    モデルの作成時とトレーニング時に使用するパラメータ セット。パラメータは、入力シェイプ、モデル構成、その他の設定を定義します。
    max_length
    データセットのサンプルの最大長。
    decode_batch_size
    Cloud TPU での Transformer 自動回帰デコードに使用されるグローバル バッチサイズ。
    decode_max_length
    デコードまたは評価データの最大シーケンス長。Cloud TPU での Transformer 自動回帰デコードで、必要なデータ パディングの量を最小限にするために使用されます。
    padded_decode
    decode_max_length にパディングされた入力データで自動回帰デコードを実行するかどうかを指定します。TPU/XLA-GPU を実行する場合、静的な形状の要件により、このフラグを設定する必要があります。
    distribution_strategy
    Cloud TPU で ResNet モデルをトレーニングするには、distribution_strategytpu に設定します。

デフォルトでは、モデルは 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. 使用する Pod スライスを指定するための tpu-sizeパラメータを使用して、ctpu up コマンドを実行します。たとえば、次のコマンドは v2-32 Pod スライスを使用します。

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

    コマンドフラグの説明

    tpu-only
    Cloud TPU のみを作成します。デフォルトでは、ctpu up コマンドは VM と Cloud TPU を作成します。
    tpu-size
    Cloud TPU のタイプ(v3-8 など)を指定します。
    zone
    Cloud TPU を作成するゾーン。これは、Compute Engine VM に使用したのと同じゾーンである必要があります。例: europe-west4-a
    tf-version
    Tensorflow ctpu のバージョンが VM にインストールされます。
    name
    作成する Cloud TPU の名前。

    ctpu ユーティリティの詳細については、CTPU リファレンスをご覧ください。

    gcloud compute ssh transformer-tutorial --zone=europe-west4-a
    
  3. TPU 名をエクスポートします。

    (vm)$ export TPU_NAME=transformer-tutorial
    
  4. モデルディレクトリの変数をエクスポートします。

    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/transformer/model_${PARAM_SET}_pod
    
  5. ディレクトリをトレーニング ディレクトリに変更します。

    (vm)$ cd /usr/share/models/official/nlp/transformer
    
  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
    

    コマンドフラグの説明

    tpu
    Cloud TPU の名前。環境変数(TPU_NAME)を指定することで設定されます。
    model_dir
    モデルのトレーニング中にチェックポイントとサマリーが保存されるディレクトリ。フォルダがない場合は、プログラムによって作成されます。Cloud TPU を使用する場合、model_dir は Cloud Storage パスにする必要があります(`gs://...`)。以前のチェックポイントが、同じサイズの TPU と TensorFlow バージョンを使用して作成されていれば、既存のフォルダを再利用して現在のチェックポイント データを読み込んで追加のチェックポイントを保存できます。
    data_dir
    トレーニング入力の Cloud Storage パス。この例では、fake_imagenet データセットに設定されています。
    vocab_file
    翻訳対象の用語が含まれるファイル。
    bleu_source
    翻訳対象のソース文が含まれるファイル。
    bleu_ref
    翻訳文のリファレンスが含まれるファイル。
    batch_size
    トレーニングのバッチサイズ。
    train_steps
    モデルをトレーニングするためのステップの数。1 つのステップで 1 つのデータバッチを処理します。これには、バックワード パスと逆伝播の両方が含まれます。
    static_batch
    データセット内のバッチに静的な形状があるかどうかを指定します。
    use_ctl
    スクリプトをカスタム トレーニング ループで実行するかどうかを指定します。
    param_set
    モデルの作成時とトレーニング時に使用するパラメータ セット。パラメータは、入力シェイプ、モデル構成、その他の設定を定義します。
    max_length
    データセットのサンプルの最大長。
    decode_batch_size
    Cloud TPU での Transformer 自動回帰デコードに使用されるグローバル バッチサイズ。
    decode_max_length
    デコードまたは評価データの最大シーケンス長。Cloud TPU での Transformer 自動回帰デコードで、必要なデータ パディングの量を最小限にするために使用されます。
    padded_decode
    decode_max_length にパディングされた入力データで自動回帰デコードを実行するかどうかを指定します。TPU/XLA-GPU を実行する場合、静的な形状の要件により、このフラグを設定する必要があります。
    steps_between_evals
    評価間に実行するトレーニング ステップの数。
    distribution_strategy
    TPU で ResNet モデルをトレーニングするには、distribution_strategytpu に設定します。

このトレーニングスクリプトは、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. 次に示すように gsutil を実行します。bucket-name の部分は、このチュートリアルで作成した Cloud Storage バケット名に置き換えてください。

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

次のステップ

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

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

ハイパーパラメータ調整

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

推論

モデルをトレーニングしたら、そのモデルを推論(予測)に使用できます。AI Platform は、機械学習モデルを開発、トレーニングデプロイするためのクラウドベースのソリューションです。モデルをデプロイすれば、AI Platform Prediction サービスを使用できるようになります。