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}
    

    このコマンドを新しい Cloud Shell VM で初めて実行すると、Authorize Cloud Shell ページが表示されます。ページの下部にある [Authorize] をクリックして、gcloud に認証情報を使用した GCP 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 europe-west4 -b on gs://bucket-name
    

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

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

    TPU VM

    $ gcloud alpha compute tpus tpu-vm create transformer-tutorial \
     --zone=europe-west4-a \
     --accelerator-type=v3-8 \
     --version=v2-alpha
    

    コマンドフラグの説明

    zone
    Cloud TPU を作成するゾーン
    accelerator-type
    作成する Cloud TPU のタイプ
    version
    Cloud TPU ランタイムのバージョン。

    TPU ノード

    $ gcloud compute tpus execution-groups create  \
    --disk-size=300 \
    --machine-type=n1-standard-8 \
    --zone=europe-west4-a \
    --tf-version=2.6.0
    

    コマンドフラグの説明

    disk-size
    VM のディスクのサイズ(GB)。
    machine_type
    gcloud コマンドで作成される VM のマシンタイプ
    tf-version
    Tensorflow gcloud compute tpus execution-groups のバージョンが VM にインストールされます。

    gcloud コマンドの詳細については、gcloud リファレンスをご覧ください。

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

    TPU VM

    gcloud alpha compute tpus tpu-vm ssh transformer-tutorial --zone=europe-west4-a
    

    TPU ノード

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

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

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

VM で、次のようにします。

  1. TensorFlow の要件をインストールし、PYTHONPATH 環境変数を設定します。

    TPU VM

    (vm)$ git clone https://github.com/tensorflow/models.git
    (vm)$ pip3 install -r models/official/requirements.txt
    
    (vm)$ export PYTHONPATH="${PWD}/models:${PYTHONPATH}"
    

    TPU ノード

    (vm)$ export PYTHONPATH="${PYTHONPATH}:/usr/share/models"
    (vm)$ pip3 install -r /usr/share/models/official/requirements.txt
    
  2. 次の環境変数を作成します。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 DATA_DIR=${HOME}/transformer/data
    (vm)$ export VOCAB_FILE=${DATA_DIR}/vocab.ende.32768
    
  3. モデルが保存されているディレクトリに移動します。

    TPU VM

    (vm)$ cd ~/models/official/nlp/transformer
    

    TPU ノード

    (vm)$ cd /usr/share/models/official/nlp/transformer
    
  4. 次のデータセットの環境変数を設定します。

    (vm)$ export GCS_DATA_DIR=${STORAGE_BUCKET}/data/transformer
    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/transformer/model_${PARAM_SET}
    
  5. データセットをダウンロードして準備します。

    (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 です。このスクリプトは 460,000 ケースを保存し、実行には約 40 分かかります。

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

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

  1. Cloud TPU 名の変数を設定します。

    TPU VM

    (vm)$ export TPU_NAME=local
    

    TPU ノード

    (vm)$ export TPU_NAME=transformer-tutorial
    
  2. トレーニング スクリプトを実行します。

    (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 Storage バケット。既存のフォルダを使用して、同じサイズの 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/tmpej76vasn
    I0218 20:07:26.020797 140707963950912 translate.py:184] Writing to file /tmp/tmpej76vasn
    I0218 20:07:35.099256 140707963950912 transformer_main.py:118] Bleu score (uncased): 0.99971704185009
    I0218 20:07:35.099616 140707963950912 transformer_main.py:119] Bleu score (cased): 0.9768599644303322
    

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

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

    (vm)$ exit
    

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

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

    TPU VM

    $ gcloud alpha compute tpus tpu-vm delete transformer-tutorial \
    --zone=europe-west4-a
    

    コマンドフラグの説明

    zone
    Cloud TPU が存在するゾーン

    TPU ノード

    $ gcloud compute tpus execution-groups delete transformer-tutorial \
    --tpu-only \
    --zone=europe-west4-a
    

    コマンドフラグの説明

    tpu-only
    Cloud TPU のみを削除します。VM は引き続き使用できます。
    zone
    削除する TPU を含むゾーン

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

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

TPU VM Pod のトレーニング

Pod プロジェクトのリソースを設定する

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

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

    Cloud Shell を開く

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

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

    gcloud config set project ${PROJECT_ID}
    

    このコマンドを新しい Cloud Shell VM で初めて実行すると、Authorize Cloud Shell ページが表示されます。ページの下部にある [Authorize] をクリックして、gcloud に認証情報を使用した GCP 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 europe-west4 gs://bucket-name
    

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

    バケットのロケーションは、VM リソースおよび TPU リソースと同じリージョンにする必要があります。

TPU VM リソースを起動する

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

    $ gcloud alpha compute tpus tpu-vm create transformer-tutorial \
    --zone=europe-west4-a \
    --accelerator-type=v3-32 \
    --version=v2-alpha-pod
    

    コマンドフラグの説明

    zone
    Cloud TPU を作成するゾーン
    accelerator-type
    作成する Cloud TPU のタイプ
    version
    Cloud TPU ランタイムのバージョン。
  2. gcloud alpha compute tpus コマンドの実行が終了したら、シェル プロンプトが username@projectname から username@vm-name に変更されたことを確認します。変更されていれば、Compute Engine VM にログインしていることになります。

    次のコマンドを実行して、Compute Engine インスタンスに接続します。

      gcloud alpha compute tpus tpu-vm ssh transformer-tutorial --zone=europe-west4-a
      

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

VM で、次のようにします。

  1. TensorFlow の要件をインストールし、PYTHONPATH 環境変数を設定します。

    (vm)$ git clone https://github.com/tensorflow/models.git
    (vm)$ pip3 install -r models/official/requirements.txt
    
    (vm)$ export PYTHONPATH="${PWD}/models:${PYTHONPATH}"
    
  2. TPU 名をエクスポートします。

    (vm)$ export TPU_NAME=transformer-tutorial
    
  3. 次の環境変数を設定します。bucket-name を Cloud Storage バケットの名前に置き換えます。

    (vm)$ export STORAGE_BUCKET=gs://bucket-name
    
    (vm)$ export TPU_LOAD_LIBRARY=0
    (vm)$ export GCS_DATA_DIR=${STORAGE_BUCKET}/data/transformer
    (vm)$ export PARAM_SET=big
    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/transformer/model_${PARAM_SET}_pod
    

    トレーニング アプリケーションでは、Cloud Storage でトレーニング データにアクセスできる必要があります。また、トレーニング アプリケーションは、Cloud Storage バケットを使用してトレーニング中にチェックポイントを保存します。

  4. スクリプト ディレクトリに移動します。

    (vm)$ cd ~/models/official/nlp/transformer
    
  5. データセットをダウンロードして準備します。

    (vm)$ export DATA_DIR=${HOME}/transformer/data
    (vm)$ export VOCAB_FILE=${DATA_DIR}/vocab.ende.32768
    (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 です。このスクリプトは 460,000 ケースを保存し、実行には約 40 分かかります。

追加のパッケージをインストールする

  1. トレーニング スクリプトには、追加のパッケージが必要です。この時点でインストールしておきます。

    (vm)$ sudo pip3 install tensorflow-model-optimization>=0.1.3
    

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

  1. 必要な環境変数を次のように設定します。

    (vm)$ export TPU_LOAD_LIBRARY=0
    (vm)$ export PYTHONPATH="${PWD}/models:${PYTHONPATH}"
    (vm)$ export RESNET_PRETRAIN_DIR=gs://cloud-tpu-checkpoints/retinanet/resnet50-checkpoint-2018-02-07
    (vm)$ export TRAIN_FILE_PATTERN=${DATA_DIR}/train-*
    (vm)$ export EVAL_FILE_PATTERN=${DATA_DIR}/val-*
    (vm)$ export VAL_JSON_FILE=${DATA_DIR}/instances_val2017.json
    
  2. トレーニング スクリプトを実行します。

    (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 Storage バケット。既存のフォルダを使用して、同じサイズの 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-32 Cloud TPU 上で約 7 分かかります。トレーニングと評価が完了すると、次のようなメッセージが表示されます。

I0415 00:28:33.108577 140097002981184 transformer_main.py:311] Train Step: 2000/2000 / loss = 5.139615058898926
I0415 00:28:33.108953 140097002981184 keras_utils.py:148] TimeHistory: 120.39 seconds, 102065.86 examples/second between steps 0 and 2000
.
.
.
I0415 00:32:01.785520 140097002981184 transformer_main.py:116] Bleu score (uncased): 0.8316259831190109
I0415 00:32:01.786150 140097002981184 transformer_main.py:117] Bleu score (cased): 0.7945530116558075

TPU ノード Pod のトレーニング

  1. 使用する Pod スライスを指定するための accelerator-type パラメータを使用して、gcloud compute tpus コマンドを実行します。たとえば、次のコマンドは v3-32 の Pod スライスを使用します。

    (vm)$ gcloud compute tpus execution-groups create \
      --tpu-only \
      --accelerator-type=v3-32 \
      --zone=europe-west4-a \
      --tf-version=2.6.0
    

    コマンドフラグの説明

    tpu-only
    Cloud TPU のみを作成します。デフォルトでは、gcloud compute tpus execution-groups コマンドは VM と Cloud TPU を作成します。
    accelerator-type
    作成する Cloud TPU のタイプ
    zone
    Cloud TPU を作成するゾーン。これは、Compute Engine VM に使用したのと同じゾーンである必要があります。例: europe-west4-a
    tf-version
    Tensorflow gcloud のバージョンが VM にインストールされます。
    name
    作成する Cloud TPU の名前。
  2. 自動的に Compute Engine インスタンスにログインしない場合は、次の ssh コマンドを実行してログインします。VM にログインすると、シェル プロンプトが username@projectname から username@vm-name に変わります。

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

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

    (vm)$ export TPU_NAME=transformer-tutorial
    
  4. 次の環境変数を設定します。bucket-name を Cloud Storage バケットの名前に置き換えます。

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

    トレーニング アプリケーションでは、Cloud Storage でトレーニング データにアクセスできる必要があります。また、トレーニング アプリケーションは、Cloud Storage バケットを使用してトレーニング中にチェックポイントを保存します。

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

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

    (vm)$ cd /usr/share/models/official/nlp/transformer
    
  7. 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 Storage バケット。既存のフォルダを使用して、同じサイズの 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 ステップごとに評価を実行します。このトレーニングと評価は、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 リソースを削除します。リソースの削除に使用するコマンドは、TPU VM または TPU ノードのどちらを使用するかによって異なります。詳細については、システム アーキテクチャをご覧ください。

    TPU VM

    $ gcloud alpha compute tpus tpu-vm delete transformer-tutorial \
    --zone=europe-west4-a
    

    TPU ノード

    $ gcloud compute tpus execution-groups delete transformer-tutorial \
    --zone=europe-west4-a
    
  3. gcloud compute tpus execution-groups list を実行して、リソースが削除されたことを確認します。削除には数分かかることがあります。以下のようなレスポンスは、インスタンスが正常に削除されたことを示します。

    $ gcloud compute tpus execution-groups list --zone=europe-west4-a
    
    Listed 0 items.
    
  4. 次に示すように 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 サービスを使用できるようになります。