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

このドキュメントでは、COCO データセットで Cloud TPU を使用して ShapeMask モデルを実行する方法を示します。

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

TPU Pod スライスでトレーニングする場合は、TPU Pod でのトレーニングを確認して、Pod スライスに必要なパラメータの変更を確認してください。

目標

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

費用

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

  • Compute Engine
  • Cloud TPU
  • クラウド ストレージ

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

始める前に

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

  1. Google Cloud アカウントにログインします。Google Cloud を初めて使用する場合は、アカウントを作成して、実際のシナリオでの Google プロダクトのパフォーマンスを評価してください。新規のお客様には、ワークロードの実行、テスト、デプロイができる無料クレジット $300 分を差し上げます。
  2. Google Cloud Console の [プロジェクト セレクタ] ページで、Google Cloud プロジェクトを選択または作成します。

    プロジェクト セレクタに移動

  3. Cloud プロジェクトに対して課金が有効になっていることを確認します。プロジェクトに対して課金が有効になっていることを確認する方法を学習する

  4. このチュートリアルでは、Google Cloud の課金対象となるコンポーネントを使用します。費用を見積もるには、Cloud TPU の料金ページを確認してください。不要な課金を回避するために、このチュートリアルを完了したら、作成したリソースを必ずクリーンアップしてください。

TPU Pod スライスでトレーニングする場合は、TPU Pod でのトレーニングを確認して、ポッドスライスに必要なパラメータの変更を確認してください。

リソースを設定する

このセクションでは、このチュートリアルで使用する Cloud Storage、VM、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 -b on gs://bucket-name
    

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

    バケットのロケーションは、仮想マシン(VM)および TPU ノードと同じリージョンにする必要があります。VM と TPU ノードは、リージョン内のサブディビジョンである特定のゾーンに配置されます。

  6. Compute Engine VM インスタンスを起動します。

    $ gcloud compute tpus execution-groups create \
     --vm-only \
     --name=shapemask-tutorial \
     --zone=europe-west4-a \
     --disk-size=300 \
     --machine-type=n1-standard-16 \
     --tf-version=2.5.0
    

    コマンドフラグの説明

    vm-only
    VM のみを作成します。デフォルトでは、gcloud compute tpus execution-groups コマンドは VM と Cloud TPU を作成します。
    name
    作成する Cloud TPU の名前。
    zone
    Cloud TPU を作成するゾーン
    disk-size
    gcloud compute tpus execution-groups コマンドで作成された VM のハードディスクのサイズ(GB)。
    machine-type
    作成する Compute Engine VM のマシンタイプ
    tf-version
    Tensorflow gcloud compute tpus execution-groups のバージョンが VM にインストールされます。

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

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

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

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

    これらの手順を続行する場合は、Compute Engine インスタンスで (vm)$ で始まる各コマンドを実行します。

  8. Cloud Storage バケットのロケーションを保存する環境変数を作成します。

    (vm)$ export STORAGE_BUCKET=gs://bucket-name
    
  9. データ ディレクトリの環境変数を作成します。

    (vm)$ export DATA_DIR=${STORAGE_BUCKET}/coco
    
  10. tpu リポジトリのクローンを作成します。

    (vm)$ git clone -b shapemask https://github.com/tensorflow/tpu/
    
  11. データの前処理に必要なパッケージをインストールします。

    (vm)$ sudo apt-get install -y python3-tk && \
      pip3 install --user Cython matplotlib opencv-python-headless pyyaml Pillow && \
      pip3 install --user "git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI"
    

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

  1. download_and_preprocess_coco.sh スクリプトを実行して、COCO データセットを、トレーニング アプリケーションで想定される一連の TFRecord(*.tfrecord)に変換します。

    (vm)$ sudo bash /usr/share/tpu/tools/datasets/download_and_preprocess_coco.sh ./data/dir/coco
    

    これにより、必要なライブラリがインストールされ、前処理スクリプトが実行されます。 ローカルのデータ ディレクトリにいくつかの *.tfrecord ファイルが出力されます。

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

    (vm)$ gsutil -m cp ./data/dir/coco/*.tfrecord ${DATA_DIR}
    (vm)$ gsutil cp ./data/dir/coco/raw-data/annotations/*.json ${DATA_DIR}
    

Cloud TPU を設定して起動する

  1. gcloud コマンドを使用して Cloud TPU リソースを起動します。

    (vm)$ gcloud compute tpus execution-groups create \
     --tpu-only \
     --accelerator-type=v3-8  \
     --name=shapemask-tutorial \
     --zone=europe-west4-a \
     --tf-version=2.5.0
    

    コマンドフラグの説明

    tpu-only
    VM を作成せずに Cloud TPU を作成します。デフォルトでは、gcloud compute tpus execution-groups コマンドは VM と Cloud TPU を作成します。
    tpu-size
    作成する Cloud TPU のタイプ
    name
    作成する Cloud TPU の名前。
    zone
    Cloud TPU を作成するゾーン
    tf-version
    Tensorflow gcloud compute tpus execution-groups のバージョンが VM にインストールされます。
  2. Cloud TPU の名前の環境変数を追加します。

    (vm)$ export TPU_NAME=shapemask-tutorial
    
  3. 指定した構成が表示されます。承認する場合は y、キャンセルする場合は n を入力してください。

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

  4. Cloud TPU の名前の環境変数を追加します。

    (vm)$ export TPU_NAME=shapemask-tutorial
    

トレーニングと評価を行う

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

  1. 必要となる環境変数を追加します。

    (vm)$ export PYTHONPATH="${PYTHONPATH}:/usr/share/models"
    (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
    (vm)$ export SHAPE_PRIOR_PATH=gs://cloud-tpu-checkpoints/shapemask/kmeans_class_priors_91x20x32x32.npy
    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/shapemask
    
  2. 次のスクリプトを実行して、ShapeMask のモデルをトレーニングします。

    (vm)$ python3 /usr/share/models/official/vision/detection/main.py \
    --strategy_type=tpu \
    --tpu=${TPU_NAME} \
    --model_dir=${MODEL_DIR} \
    --mode=train \
    --model=shapemask \
    --params_override="{train: {total_steps: 100, learning_rate: {init_learning_rate: 0.08, learning_rate_levels: [0.008, 0.0008], learning_rate_steps: [15000, 20000], }, checkpoint: { path: ${RESNET_CHECKPOINT},prefix: resnet50}, train_file_pattern: ${TRAIN_FILE_PATTERN}}, shapemask_head: {use_category_for_mask: true, shape_prior_path: ${SHAPE_PRIOR_PATH}}, shapemask_parser: {output_size: [640, 640]}}"
    

    コマンドフラグの説明

    strategy_type
    TPU で RetinaNet モデルをトレーニングするには、distribution_strategytpu に設定する必要があります。
    tpu
    Cloud TPU の名前。TPU_NAME 環境変数を使用して設定します。
    model_dir
    トレーニング中にチェックポイントとサマリーが保存される Cloud Storage バケット。既存のフォルダを使用して、同じサイズの TPU と TensorFlow バージョンで以前に生成されたチェックポイントを読み込むことができます。
    mode
    trainevaltrain_and_eval のいずれか。
    model
    トレーニングするモデル。
    params_override
    デフォルトのスクリプト パラメータをオーバーライドする JSON 文字列。スクリプト パラメータの詳細については、/usr/share/models/official/vision/detection/main.py をご覧ください。

    トレーニング スクリプトの出力は次のようになります。

    Train Step: 100/100  / loss = {
     'total_loss': 10.152124404907227,
     'loss': 10.152124404907227,
     'retinanet_cls_loss': 1.3409545421600342,
     'l2_regularization_loss': 4.6183762550354,
     'retinanet_box_loss': 0.012389584444463253,
     'shapemask_prior_loss': 0.183448925614357,
     'shapemask_coarse_mask_loss': 1.7648673057556152,
     'shapemask_fine_mask_loss': 1.790102243423462,
     'model_loss': 5.533748626708984,
     'learning_rate': 0.021359999
    }
    / training metric = {
    'total_loss': 10.152124404907227,
    'loss': 10.152124404907227,
    'retinanet_cls_loss': 1.3409545421600342,
    'l2_regularization_loss': 4.6183762550354,
    'retinanet_box_loss': 0.012389584444463253,
    'shapemask_prior_loss': 0.183448925614357,
    'shapemask_coarse_mask_loss': 1.7648673057556152,
    'shapemask_fine_mask_loss': 1.790102243423462,
    'model_loss': 5.533748626708984,
    'learning_rate': 0.021359999
    }
  3. ShapeMask モデルを評価するスクリプトを実行します。

    (vm)$ python3 /usr/share/models/official/vision/detection/main.py \
    --strategy_type=tpu \
    --tpu=${TPU_NAME} \
    --model_dir=${MODEL_DIR} \
    --checkpoint_path=${MODEL_DIR} \
    --mode=eval_once \
    --model=shapemask \
    --params_override="{eval: { val_json_file: ${VAL_JSON_FILE}, eval_file_pattern: ${EVAL_FILE_PATTERN}, eval_samples: 5000 }, shapemask_head: {use_category_for_mask: true, shape_prior_path: ${SHAPE_PRIOR_PATH}}, shapemask_parser: {output_size: [640, 640]}}"
    

    コマンドフラグの説明

    strategy_type
    TPU で RetinaNet モデルをトレーニングするには、distribution_strategytpu に設定する必要があります。
    tpu
    Cloud TPU の名前。TPU_NAME 環境変数を使用して設定します。
    model_dir
    トレーニング中にチェックポイントとサマリーが保存される Cloud Storage バケット。既存のフォルダを使用して、同じサイズの TPU と TensorFlow バージョンで以前に生成されたチェックポイントを読み込むことができます。
    mode
    trainevaltrain_and_eval のいずれか。
    model
    トレーニングするモデル。
    params_override
    デフォルトのスクリプト パラメータをオーバーライドする JSON 文字列。スクリプト パラメータの詳細については、/usr/share/models/official/vision/detection/main.py をご覧ください。

評価スクリプトの出力は次のようになります。

756s 62ms/step
- loss: 0.4864
- accuracy: 0.8055
- val_loss: 0.3832
- val_accuracy: 0.8546

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

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

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

  • v2-32
  • v3-32

Cloud TPU Pod を使用する場合は、まず Pod を使用してモデルをトレーニングし、その後、単一の Cloud TPU デバイスを使用してモデルを評価します。

Cloud TPU Pod を使用したトレーニング

Compute Engine インスタンスをすでに削除している場合は、リソースのセットアップの手順に沿って新しいインスタンスを作成します。

以下のサンプルトレーニングは、わずか 20 ステップで行われ、v3-32 TPU ノードで完了するまでに約 10 分かかります。収束するようにトレーニングするには、v3-32 TPU Pod では約 11,250 ステップ、約 2 時間かかります。

  1. 単一のデバイスでモデルをトレーニングするために作成した Cloud TPU リソースを削除します。

    (vm)$ gcloud compute tpus execution-groups delete shapemask-tutorial \
      --zone=europe-west4-a \
      --tpu-only
  2. 使用する Pod スライスを指定するための accelerator-typeパラメータを使用して、gcloud compute tpus execution-groups コマンドを実行します。たとえば、次のコマンドは v3-32 の Pod スライスを使用します。

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

    コマンドフラグの説明

    name
    作成する Cloud TPU の名前。
    accelerator-type
    作成する Cloud TPU のタイプ
    zone
    Cloud TPU を作成するゾーン
    tf-version
    Tensorflow gcloud のバージョンが VM にインストールされます。
    tpu-only
    Cloud TPU のみを作成します。デフォルトでは、gcloud コマンドは VM と Cloud TPU を作成します。
  3. TPU_NAME と MODEL_DIR 環境変数を更新します。

    (vm)$ export TPU_NAME=shapemask-tutorial
    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/shapemask-pods
    
  4. トレーニング スクリプトを起動します。

    (vm)$ python3 /usr/share/models/official/vision/detection/main.py \
    --strategy_type=tpu \
    --tpu=${TPU_NAME} \
    --model_dir=${MODEL_DIR} \
    --mode=train \
    --model=shapemask \
    --params_override="{train: { batch_size: 128, iterations_per_loop: 500, total_steps: 20, learning_rate: {'learning_rate_levels': [0.008, 0.0008],'learning_rate_steps': [10000, 13000]}, checkpoint: {path: ${RESNET_CHECKPOINT}, prefix: resnet50/}, train_file_pattern: ${TRAIN_FILE_PATTERN}}, eval: {val_json_file: ${VAL_JSON_FILE}, eval_file_pattern: ${EVAL_FILE_PATTERN}},shapemask_head: {use_category_for_mask: true, shape_prior_path: ${SHAPE_PRIOR_PATH}}}"
    

    コマンドフラグの説明

    strategy_type
    TPU で RetinaNet モデルをトレーニングするには、distribution_strategytpu に設定する必要があります。
    tpu
    Cloud TPU の名前。TPU_NAME 環境変数を使用して設定します。
    model_dir
    トレーニング中にチェックポイントとサマリーが保存される Cloud Storage バケット。既存のフォルダを使用して、同じサイズの TPU と TensorFlow バージョンで以前に生成されたチェックポイントを読み込むことができます。
    mode
    trainevaltrain_and_eval のいずれか。
    model
    トレーニングするモデル。
    params_override
    デフォルトのスクリプト パラメータをオーバーライドする JSON 文字列。スクリプト パラメータの詳細については、/usr/share/models/official/vision/detection/main.py をご覧ください。

トレーニング スクリプトの出力は次のようになります。

Train Step: 20/20  / loss = {
  'total_loss': 12.213006973266602,
  'loss': 12.213006973266602,
  'retinanet_cls_loss': 1.9299328327178955,
  'l2_regularization_loss': 4.628948211669922,
  'retinanet_box_loss': 0.016126759350299835,
  'shapemask_prior_loss': 0.16990719735622406,
  'shapemask_coarse_mask_loss': 3.688129425048828,
  'shapemask_fine_mask_loss': 1.1426670551300049,
  'model_loss': 7.584057807922363,
  'learning_rate': 0.009632
}
/ training metric = {
  'total_loss': 12.213006973266602,
  'loss': 12.213006973266602,
  'retinanet_cls_loss': 1.9299328327178955,
  'l2_regularization_loss': 4.628948211669922,
  'retinanet_box_loss': 0.016126759350299835,
  'shapemask_prior_loss': 0.16990719735622406,
  'shapemask_coarse_mask_loss': 3.688129425048828,
  'shapemask_fine_mask_loss': 1.1426670551300049,
  'model_loss': 7.584057807922363,
  'learning_rate': 0.009632
}

モデルの評価

このステップでは、単一の Cloud TPU ノードを使用して、COCO データセットに対して上記でトレーニングしたモデルを評価します。評価には約 20 分かかります。

  1. ポッドでモデルをトレーニングするために作成した Cloud TPU リソースを削除します。

    (vm)$ gcloud compute tpus execution-groups delete shapemask-tutorial \
      --zone=europe-west4-a \
      --tpu-only
  2. 評価を実行するために、新しい TPU デバイスを起動します。

    (vm)$ gcloud compute tpus execution-groups create --tpu-only \
      --name=shapemask-tutorial \
      --accelerator-type=v3-8 \
      --zone=europe-west4-a \
      --tf-version=2.5.0 \
    
    

    コマンドフラグの説明

    tpu-only
    Cloud TPU のみを作成します。デフォルトでは、gcloud コマンドは VM と Cloud TPU を作成します。
    name
    作成する Cloud TPU の名前。
    accelerator-type
    作成する Cloud TPU のタイプ
    zone
    Cloud TPU を作成するゾーン
    tf-version
    Tensorflow gcloud のバージョンが VM にインストールされます。
  3. TPU_NAME 環境変数を更新します。

    (vm)$ export TPU_NAME=shapemask-tutorial
    
  4. 評価を開始します。

    (vm)$ python3 /usr/share/models/official/vision/detection/main.py \
    --strategy_type=tpu \
    --tpu=shapemask-tutorial \
    --model_dir=${MODEL_DIR} \
    --checkpoint_path=${MODEL_DIR} \
    --mode=eval_once \
    --model=shapemask \
    --params_override="{eval: { val_json_file: ${VAL_JSON_FILE}, eval_file_pattern: ${EVAL_FILE_PATTERN}, eval_samples: 5000 } }"
    

    コマンドフラグの説明

    strategy_type
    TPU で RetinaNet モデルをトレーニングするには、distribution_strategytpu に設定する必要があります。
    tpu
    Cloud TPU の名前。TPU_NAME 環境変数を使用して設定します。
    model_dir
    トレーニング中にチェックポイントとサマリーが保存される Cloud Storage バケット。既存のフォルダを使用して、同じサイズの TPU と TensorFlow バージョンで以前に生成されたチェックポイントを読み込むことができます。
    mode
    trainevaltrain_and_eval のいずれか。
    model
    トレーニングするモデル。
    params_override
    デフォルトのスクリプト パラメータをオーバーライドする JSON 文字列。スクリプト パラメータの詳細については、/usr/share/models/official/vision/detection/main.py をご覧ください。

クリーンアップ

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

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

    (vm)$ exit
    

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

  2. Cloud Shell で次のコマンドを使用して、Compute Engine VM と Cloud TPU を削除します。

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

    $ gcloud compute tpus execution-groups list \
     --zone=europe-west4-a
    

    次のような空の TPU のリストが表示されます。

       NAME             STATUS
    
  4. 次に示すように、gsutil を使用して Cloud Storage バケットを削除します。bucket-name を Cloud Storage バケットの名前に置き換えます。

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

次のステップ

さまざまなサイズの画像でトレーニングする

より規模の大きいニューラル ネットワーク(たとえば ResNet-50 ではなく ResNet-101)を使用して調査できます。入力画像のサイズを大きくして、より強力なニューラル ネットワークを使用すると、処理時間は長くなりますが、モデルの精度が高くなります。

別のベースを使用する

また、独自のデータセットで ResNet モデルを事前トレーニングし、それを ShapeMask モデルのベースとして使用することもできます。いくつかの作業を行うことで、ResNet の代わりに別のニューラル ネットワークに切り替えることもできます。最後に、独自のオブジェクト検出モデルの実装を検討されている場合は、このネットワークを今後のテストに利用することをおすすめします。