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


概要

このチュートリアルでは、COCO データセットで Cloud TPU を使用して Mask RCNN モデルを実行する方法を示します。

Mask RCNN は、コンピュータ ビジョンの難しい課題の一つであるオブジェクト検出と画像セグメンテーションに対応するように設計されたディープ ニューラル ネットワークです。

Mask RCNN モデルは、画像内の個々のオブジェクトそれぞれのインスタンスに対し、境界ボックスとセグメンテーション マスクを生成します。このモデルは、Feature Pyramid Network(FPN)および ResNet50 バックボーンに基づいています。

このチュートリアルでは、Tensorflow Keras API を使用してモデルをトレーニングします。Keras API は上位の TensorFlow API であり、Cloud TPU で機械学習モデルを作成して実行するのに使用できます。この API を使用すると、低レベルの実装の大部分が隠されることで、モデル開発プロセスが簡素化され、TPU と他のプラットフォーム(GPU や CPU など)の切り替えが簡単になります。

以下の手順では、Cloud TPU でモデルをトレーニングする方法をすでに理解していることを前提としています。Cloud TPU を初めて使用する場合は、クイックスタートで基本的な概要をご確認ください。

目標

  • COCO データセットを準備する
  • データセットとモデルの出力を格納する Cloud Storage バケットを作成する
  • トレーニングと評価のための TPU リソースを設定する
  • 単一の Cloud TPU または Cloud TPU Pod でトレーニングと評価を実行する

費用

このドキュメントでは、Google Cloud の次の課金対象のコンポーネントを使用します。

  • Compute Engine
  • Cloud TPU
  • Cloud Storage

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

始める前に

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

  1. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.
  2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  3. Make sure that billing is enabled for your Google Cloud project.

  4. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  5. Make sure that billing is enabled for your Google Cloud project.

  6. このチュートリアルでは、Google Cloud の課金対象となるコンポーネントを使用します。費用を見積もるには、Cloud TPU の料金ページを確認してください。

COCO データセットを準備する

このチュートリアルでは、COCO データセットを使用します。トレーニングに使用するデータセットは、Cloud Storage バケットの TFRecord 形式である必要があります。

モデルのトレーニングに使用するゾーンにある Cloud Storage バケットで、すでに COCO データセットを準備している場合、直接単一デバイスのトレーニングに移行できます。それ以外の場合は、次の手順でデータセットを準備します。

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

    Cloud Shell を開く

  2. Cloud Shell で、gcloud をプロジェクト ID で構成します。

    export PROJECT_ID=project-id
    gcloud config set project ${PROJECT_ID}
  3. Cloud Shell で、次のコマンドを使用して Cloud Storage バケットを作成します。

    gcloud storage buckets create gs://bucket-name --project=${PROJECT_ID} --location=us-central2
  4. データセットをダウンロードして前処理するための Compute Engine VM を作成します。詳細については、Compute Engine インスタンスを作成して起動するをご覧ください。

    $ gcloud compute instances create vm-name \
        --zone=us-central2-b \
        --image-family=ubuntu-2204-lts \
        --image-project=ubuntu-os-cloud \
        --machine-type=n1-standard-16 \
        --boot-disk-size=300GB
  5. SSH を使用して Compute Engine VM に接続します。

    $ gcloud compute ssh vm-name --zone=us-central2-b

    VM に接続すると、シェル プロンプトが username@projectname から username@vm-name に変わります。

  6. 2 つの変数を設定します。1 つは以前に作成したストレージ バケット用であり、もう 1 つはストレージ バケット上のトレーニング データ(DATA_DIR)を保持するディレクトリ用です。

    (vm)$ export STORAGE_BUCKET=gs://bucket-name
    (vm)$ export DATA_DIR=${STORAGE_BUCKET}/coco
  7. データの前処理に必要なパッケージをインストールします。

    (vm)$ sudo apt-get update && \
      sudo apt-get install python3-pip && \
      sudo apt-get install -y python3-tk && \
      pip3 install --user Cython matplotlib opencv-python-headless pyyaml Pillow numpy absl-py tensorflow && \
      pip3 install --user "git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI"
  8. download_and_preprocess_coco.sh スクリプトを実行して、COCO データセットを、トレーニング アプリケーションで想定される一連の TFRecord ファイル(*.tfrecord)に変換します。

    (vm)$ git clone https://github.com/tensorflow/tpu.git
    (vm)$ sudo bash tpu/tools/datasets/download_and_preprocess_coco.sh ./data/dir/coco

    これにより、必要なライブラリがインストールされ、前処理スクリプトが実行されます。ローカルのデータ ディレクトリに *.tfrecord ファイルが出力されます。COCO のダウンロードと変換スクリプトが完了するまでには約 1 時間かかります。

  9. データを Cloud Storage バケットにコピーする

    データを TFRecord 形式に変換してから、gcloud CLI を使用してローカル ストレージから Cloud Storage バケットにデータをコピーします。アノテーション ファイルもコピーする必要があります。アノテーション ファイルは、モデルのパフォーマンスの検証に利用できます。

    (vm)$ gcloud storage cp ./data/dir/coco/*.tfrecord ${DATA_DIR}
    (vm)$ gcloud storage cp ./data/dir/coco/raw-data/annotations/*.json ${DATA_DIR}
  10. Compute Engine VM との接続を解除します。

    (vm)$ exit

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

  11. Compute Engine VM を削除します。

    $ gcloud compute instances delete vm-name \
    --zone=us-central2-b

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}
 

クリーンアップ

このチュートリアルで使用したリソースについて、Google Cloud アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。

トレーニングを実施した後、TPU VM を削除し、ストレージ バケットを削除します。

  1. Compute Engine インスタンスとの接続を切断していない場合は切断します。

    (vm)$ exit

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

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

    $ gcloud compute tpus tpu-vm delete mask-rcnn-tutorial \
    --zone=${ZONE}
  3. gcloud compute tpus tpu-vm list を実行して、リソースが削除されたことを確認します。削除には数分かかることがあります。次のコマンドの出力には、このチュートリアルで作成したリソースを含めないでください。

    $ gcloud compute tpus tpu-vm list --zone=${ZONE}
  4. 次に示すように gcloud CLI を実行します。bucket-name は、このチュートリアルで作成した Cloud Storage バケット名に置き換えます。

    $ gcloud storage rm gs://bucket-name --recursive

次のステップ

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 推論の概要をご覧ください。

TensorBoard の TPU ツールを確認する。