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

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

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 に認証情報を使用した Google Cloud 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. TPU の設定変数をエクスポートする

   プロジェクト ID、TPU リソースに使用する名前、モデルをトレーニングしてトレーニング関連のデータを保存するためのゾーンをエクスポートします。

   $ export TPU_NAME=mask-rcnn-tutorial
   $ export ZONE=europe-west4-a

6. Compute Engine VM と Cloud TPU を起動します。

   $ gcloud compute tpus tpu-vm create mask-rcnn-tutorial \
      --zone=${ZONE} \
      --accelerator-type=v3-8 \
      --version=tpu-vm-tf-2.17.0-pjrt
     

   コマンドフラグの説明

   zone

   Cloud TPU を作成するゾーン。

   accelerator-type

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

   version

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

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

7. SSH を使用して Compute Engine インスタンスに接続します。VM に接続すると、シェル プロンプトが username@projectname から username@vm-name に変わります。

   gcloud compute tpus tpu-vm ssh mask-rcnn-tutorial --zone=${ZONE}

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

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

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

   (vm)$ export TPU_NAME=local

10. 次の環境変数を設定します。bucket-name は、COCO データセットを格納する Cloud Storage バケットの名前に置き換えます。

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

11. データ ディレクトリとモデル ディレクトリの環境変数を追加します。

    (vm)$ export DATA_DIR=${STORAGE_BUCKET}/coco
    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/mask-rcnn

12. 他に必要な環境変数を追加します。

    (vm)$ export RESNET_CHECKPOINT=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

13. PYTHONPATH 環境変数を設定します。

    (vm)$ export PYTHONPATH="${PYTHONPATH}:/usr/share/tpu/models"

14. TPU を作成するときに、--version パラメータを -pjrt で終わるバージョンに設定した場合は、次の環境変数を設定して PJRT ランタイムを有効にします。

      (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
      (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

15. モデルが保存されているディレクトリに移動します。

    (vm)$ cd /usr/share/tpu/models/official/vision

次のスクリプトは、10 のトレーニング ステップと 10 の評価ステップでトレーニングするサンプル トレーニングを実行します。v3-8 TPU で完了するまでに約 6 分かかります。収束するようにトレーニングするには、v3-8 TPU では約 22,500 ステップ、約 6 時間かかります。

1. 次のコマンドを実行して Mask-RCNN モデルをトレーニングします。

   (vm)$ python3 train.py \
     --tpu=${TPU_NAME} \
     --experiment=maskrcnn_resnetfpn_coco \
     --mode=train_and_eval \
     --config_file=configs/experiments/maskrcnn/r50fpn_640_coco_scratch_tpu4x4.yaml \
     --model_dir=${MODEL_DIR} \
     --params_override="task.train_data.input_path=${TRAIN_FILE_PATTERN},task.validation_data.input_path=${EVAL_FILE_PATTERN},task.annotation_file=${VAL_JSON_FILE},runtime.distribution_strategy=tpu,trainer.train_steps=10,trainer.validation_steps=10,task.train_data.global_batch_size=8,task.validation_data.global_batch_size=8"

   コマンドフラグの説明

   strategy_type

   配信戦略。

   tpu

   TPU の名前。

   model_dir

   モデルのトレーニング中にチェックポイントとサマリーが保存されるディレクトリを指定します。フォルダがない場合は、プログラムによって作成されます。Cloud TPU を使用する場合、model_dir は Cloud Storage パスにする必要があります（`gs://...`）。以前のチェックポイントが、同じサイズの TPU と TensorFlow バージョンを使用して作成されていれば、既存のフォルダを再利用して現在のチェックポイント データを読み込んで追加のチェックポイントを保存できます。

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

   {'frcnn_box_loss': 0.033865165,
    'frcnn_cls_loss': 1.2535654,
    'learning_rate': 0.008266499,
    'mask_loss': 1.2039567,
    'model_loss': 2.821458,
    'rpn_box_loss': 0.034982488,
    'rpn_score_loss': 0.2950886,
    'total_loss': 4.340171,
    'training_loss': 4.340171}
   train | step:     10 | steps/sec:    0.1 | output:
   {'frcnn_box_loss': 0.033865165,
    'frcnn_cls_loss': 1.2535654,
    'learning_rate': 0.008266499,
    'mask_loss': 1.2039567,
    'model_loss': 2.821458,
    'rpn_box_loss': 0.034982488,
    'rpn_score_loss': 0.2950886,
    'total_loss': 4.340171,
    'training_loss': 4.340171}
   

   その後に、評価ステップからの出力が続きます。

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

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

   (vm)$ exit

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

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

   $ gcloud compute tpus tpu-vm delete mask-rcnn-tutorial \
       --zone=europe-west4-a

   コマンドフラグの説明

   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 に認証情報を使用した 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. 以前に COCO データセットを準備してストレージ バケットに移動した場合は、これを再び Pod トレーニングに使用できます。COCO データセットをまだ準備していない場合は、今すぐ準備し、ここに戻ってトレーニングを設定してください。

6. Cloud TPU Pod を起動する

   このチュートリアルでは、v3-32 Pod を指定します。他の Pod オプションについては、TPU のバージョンをご覧ください。

   $ gcloud compute tpus tpu-vm create mask-rcnn-tutorial \
     --zone=${ZONE} \
     --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 ソフトウェアのバージョン。

7. SSH を使用して Compute Engine インスタンスに接続します。VM に接続すると、シェル プロンプトが username@projectname から username@vm-name に変わります。

   gcloud compute tpus tpu-vm ssh mask-rcnn-tutorial --zone=${ZONE}

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

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

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

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

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

    (vm)$ export TPU_NAME=mask-rcnn-tutorial

11. 次の環境変数を設定します。bucket-name を Cloud Storage バケットの名前に置き換えます。

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

12. 他に必要な環境変数を追加します。

    (vm)$ export RESNET_CHECKPOINT=gs://cloud-tpu-checkpoints/retinanet/resnet50-checkpoint-2018-02-07
    (vm)$ export DATA_DIR=${STORAGE_BUCKET}/coco
    (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
    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/mask-rcnn-pod

13. PYTHONPATH 環境変数を設定します。

    (vm)$ export PYTHONPATH="/usr/share/tpu/models:${PYTHONPATH}"
    (vm)$ export TPU_LOAD_LIBRARY=0

14. モデルが保存されているディレクトリに移動します。

    (vm)$ cd /usr/share/tpu/models/official/vision

    (vm)$ cd /usr/share/models/official/legacy/detection

15. モデルをトレーニングします。

    この手順では、COCO データセットで 10 トレーニング ステップでモデルをトレーニングします。このトレーニングは、v3-32 Cloud TPU 上で約 10 分かかります。

    (vm)$ python3 train.py \
      --tpu=${TPU_NAME} \
      --experiment=maskrcnn_resnetfpn_coco \
      --mode=train_and_eval \
      --config_file=configs/experiments/maskrcnn/r50fpn_640_coco_scratch_tpu4x4.yaml \
      --model_dir=${MODEL_DIR} \
      --params_override="task.train_data.input_path=${TRAIN_FILE_PATTERN},task.validation_data.input_path=${EVAL_FILE_PATTERN},task.annotation_file=${VAL_JSON_FILE},runtime.distribution_strategy=tpu,trainer.train_steps=10,trainer.validation_steps=10,task.train_data.global_batch_size=256,task.validation_data.global_batch_size=256"

    コマンドフラグの説明

    tpu

    TPU の名前。

    model_dir

    モデルのトレーニング中にチェックポイントとサマリーが保存されるディレクトリを指定します。該当するフォルダがない場合は、プログラムによって作成されます。Cloud TPU を使用する場合は、model_dir を Cloud Storage パスにする必要があります（gs://...）。以前のチェックポイントが、同じサイズの Cloud TPU と TensorFlow バージョンを使用して作成されていれば、既存のフォルダを再利用して現在のチェックポイント データを読み込んで追加のチェックポイントを保存できます。

    params_override

    デフォルトのスクリプト パラメータをオーバーライドする JSON 文字列。

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

 I0706 19:47:16.108213 139955064548416 controller.py:457] train | step: 10 | steps/sec:    0.1 | output:
    {'frcnn_box_loss': 0.05632668,
     'frcnn_cls_loss': 1.3012192,
     'learning_rate': 0.008266499,
     'mask_loss': 1.2371812,
     'model_loss': 2.9746659,
     'rpn_box_loss': 0.08227444,
     'rpn_score_loss': 0.2976642,
     'total_loss': 4.493513,
     'training_loss': 4.493513}
train | step:     10 | steps/sec:    0.1 | output:
    {'frcnn_box_loss': 0.05632668,
     'frcnn_cls_loss': 1.3012192,
     'learning_rate': 0.008266499,
     'mask_loss': 1.2371812,
     'model_loss': 2.9746659,
     'rpn_box_loss': 0.08227444,
     'rpn_score_loss': 0.2976642,
     'total_loss': 4.493513,
     'training_loss': 4.493513}