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

このドキュメントでは、RetinaNet オブジェクト検出モデルの実装について説明します。コードは、GitHub で入手できます。

以下の手順では、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
  • Cloud Storage

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

始める前に

このセクションでは、Cloud Storage バケットと Compute Engine VM の設定について説明します。

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

    Cloud Shell を開く

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

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

    gcloud config set project ${PROJECT_ID}
    
  4. 次のコマンドを使用して Cloud Storage バケットを作成します。

    gsutil mb -p ${PROJECT_ID} -c standard -l europe-west4 -b on gs://bucket-name
    

    この Cloud Storage バケットには、モデルのトレーニングに使用するデータとトレーニング結果が格納されます。

  5. ctpu up コマンドを使用して、Compute Engine VM を起動します。

    $ ctpu up --vm-only \
    --name=retinanet-tutorial \
    --disk-size-gb=300 \
    --machine-type=n1-standard-8 \
    --zone=europe-west4-a \
    --tf-version=2.2
    

    コマンドフラグの説明

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

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

  6. プロンプトが表示されたら、y キーを押して Cloud TPU リソースを作成します。

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

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

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

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

    RetinaNet トレーニング アプリケーションには、いくつかの追加パッケージが必要です。 この時点で、これらのパッケージをインストールしておきます。

    (vm)$ sudo apt-get install -y python3-tk
    (vm)$ pip3 install --user Cython matplotlib opencv-python-headless pyyaml Pillow
    
    (vm)$ pip3 install --user 'git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI'
    
    (vm)$ sudo pip3 install --user -r /usr/share/models/official/requirements.txt
    

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

COCO データセットは Cloud Storage に保存されるため、作成したバケットの名前を指定してストレージ バケット変数を設定します。

(vm)$ export STORAGE_BUCKET=gs://bucket-name
(vm)$ export DATA_DIR=${STORAGE_BUCKET}/coco

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 ファイルが出力されます。COCO のダウンロードと変換スクリプトが完了するまでには約 1 時間かかります。

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

データを 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}

トレーニング環境を設定する

  1. ctpu up コマンドを使用して Cloud TPU を作成します。

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

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

    コマンドフラグの説明

    tpu-only
    Cloud TPU のみを作成します。デフォルトでは、ctpu up コマンドは VM と Cloud TPU を作成します。
    tpu-size
    作成する Cloud TPU のタイプ
    zone
    Cloud TPU を作成するゾーン
    name
    作成する Cloud TPU の名前。
    tf-version
    Tensorflow ctpu のバージョンが VM にインストールされます。
    パラメータ 説明

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

  2. プロンプトが表示されたら、y キーを押して Cloud TPU リソースを作成します。

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

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

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

  3. TPU 名の環境変数を作成します。

    (vm)$ export TPU_NAME=retinanet-tutorial
    
  4. 次のコマンドを使用して、Python パスに最上位の /models フォルダを追加します。

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

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

以下のトレーニング スクリプトは、Cloud TPU v3-8 で実行されたものです。 トレーニング時間は長くなりますが、Cloud TPU v2-8 でも実行できます。

以下のサンプルスクリプトは、わずか 10 ステップのトレーニングで、v3-8 TPU ノードでは 5 分未満で実行されます。収束するようにトレーニングするには、Cloud TPU v3-8 TPU では約 22,500 ステップ、約 1 時間半かかります。

  1. 次の環境変数を設定します。

    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/retinanet-train
    (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
    
  2. トレーニング スクリプトを実行します。

    (vm)$ python3 /usr/share/models/official/vision/detection/main.py \
         --strategy_type=tpu \
         --tpu=${TPU_NAME} \
         --model_dir=${MODEL_DIR} \
         --mode="train" \
         --params_override="{ type: retinanet, train: { total_steps: 10, 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}, eval_samples: 5000 } }"
    

    コマンドフラグの説明

    strategy_type
    TPU で RetinaNet モデルをトレーニングするには、distribution_strategytpu に設定する必要があります。
    tpu
    Cloud TPU の名前。TPU_NAME 環境変数を使用して設定します。
    model_dir
    モデルのトレーニング中にチェックポイントとサマリーが保存されるディレクトリ。フォルダがない場合は、プログラムによって作成されます。Cloud TPU を使用する場合、model_dir は Cloud Storage のパスにする必要があります(gs://...)。以前に作成されたチェックポイントの Cloud TPU のサイズと TensorFlow のバージョンが同じであれば、既存のフォルダを再利用して現在のチェックポイント データを読み込み、追加のチェックポイントを保存できます。
    mode
    モデルをトレーニングするには train、モデルを評価するには eval に設定します。
    params_override
    デフォルトのスクリプト パラメータをオーバーライドする JSON 文字列。スクリプト パラメータの詳細については、/usr/share/models/official/vision/detection/main.py をご覧ください。

モデルのトレーニング中は、ログ出力を表示して進行状況を確認できます。次のような出力は、トレーニングが正常に進行していることを示しています。

31517803669, 'learning_rate': 0.08, 'box_loss': 0.0006472870009019971,
'l2_regularization_loss': 0.09328179806470871}
I0210 21:59:19.888985 139927795508992 distributed_executor.py:49]
Saving model as TF checkpoint: gs://bucket-eu/retinanet-model/ctl_step_2500.ckpt-5
I0210 22:01:07.714749 139927795508992 distributed_executor.py:446]
Train Step: 3000/22500  / loss = {'model_loss': 0.08362223953008652,
'total_loss': 0.17120523750782013, 'cls_loss': 0.057121846824884415,
'learning_rate': 0.08, 'box_loss': 0.0005300078773871064,
'l2_regularization_loss': 0.08758299797773361} / training metric =
{'model_loss': 0.08362223953008652, 'total_loss': 0.17120523750782013, 'cls_loss': 0.0
57121846824884415, 'learning_rate': 0.08, 'box_loss': 0.0005300078773871064,
'l2_regularization_loss': 0.08758299797773361}
I0210 22:01:15.813422 139927795508992 distributed_executor.py:49]
Saving model as TF checkpoint: gs://bucket-eu/retinanet-model/ctl_step_3000.ckpt-6

単一の Cloud TPU デバイスの評価

次の手順では、COCO 評価データを使用します。評価ステップが完了するまでには約 10 分かかります。

  1. 次の環境変数を設定します。

    (vm)$ export EVAL_SAMPLES=5000
    
  2. 評価スクリプトを実行します。

    (vm)$ python3 /usr/share/models/official/vision/detection/main.py \
          --strategy_type=tpu \
          --tpu=${TPU_NAME} \
          --model_dir=${MODEL_DIR} \
          --mode="eval" \
          --params_override="{ type: retinanet, eval: { val_json_file: ${VAL_JSON_FILE}, eval_file_pattern: ${EVAL_FILE_PATTERN}, eval_samples: ${EVAL_SAMPLES} } }"
    

    コマンドフラグの説明

    strategy_type
    TPU で RetinaNet モデルをトレーニングするには、distribution_strategytpu に設定する必要があります。
    tpu
    Cloud TPU の名前。TPU_NAME 環境変数を使用して設定します。
    model_dir
    モデルのトレーニング中にチェックポイントとサマリーが保存されるディレクトリ。フォルダがない場合は、プログラムによって作成されます。Cloud TPU を使用する場合、model_dir は Cloud Storage のパスにする必要があります(gs://...)。以前に作成されたチェックポイントの Cloud TPU のサイズと TensorFlow のバージョンが同じであれば、既存のフォルダを再利用して現在のチェックポイント データを読み込み、追加のチェックポイントを保存できます。
    mode
    modeleval に設定して、モデルを評価します。
    params_override
    デフォルトのスクリプト パラメータをオーバーライドする JSON 文字列。スクリプト パラメータの詳細については、/usr/share/models/official/vision/detection/main.py をご覧ください。

    評価の最後に、コンソールに次のようなメッセージが表示されます。

    Accumulating evaluation results...
    DONE (t=7.66s).
     Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.000
     Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.000
     Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.000
     Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
     Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.000
     Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.000
     Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.000
     Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.000
     Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.000
     Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
     Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.000
     Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.000
    

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

Cloud TPU Pod を使用してモデルをスケーリングする

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

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

    (vm)$ ctpu delete --tpu-only --zone=europe-west4-a --name=retinanet-tutorial
  2. 使用する Pod スライスを指定するための tpu-size パラメータを使用して、ctpu up コマンドを実行します。たとえば、次のコマンドは v2-32 Pod スライスを使用します。

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

    コマンドフラグの説明

    tpu-only
    Cloud TPU のみを作成します。デフォルトでは、ctpu up コマンドは VM と Cloud TPU の両方を作成します。
    tpu-size
    作成する Cloud TPU のタイプ
    zone
    Cloud TPU を作成するゾーン
    name
    作成する Cloud TPU の名前。
    tf-version
    Tensorflow ctpu のバージョンが VM にインストールされます。

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

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

    (vm)$ export TPU_NAME=retinanet-tutorial
    
  4. 次の環境変数を設定します。

    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/retinanet-pod
    
  5. v2-32 TPU ノードで Pod トレーニング スクリプトを実行します。

    次のサンプル トレーニング スクリプトは、Cloud TPU v2-32 Podで実行されたものです。トレーニングはわずか 10 ステップで、5 分未満で実行されます。収束するようにトレーニングするには、v2-32 TPU Pod では 2,109 ステップ、約 50 分かかります。

    (vm)$  python3 /usr/share/models/official/vision/detection/main.py \
        --strategy_type=tpu \
        --tpu=${TPU_NAME} \
        --model_dir=${MODEL_DIR} \
        --mode="train" \
        --params_override="{ type: retinanet, train: { total_steps: 10, batch_size: 256, 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}, eval_samples: 5000 } }"
    

    コマンドフラグの説明

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

クリーンアップ

このチュートリアルで使用したリソースについて、Google Cloud Platform アカウントに課金されないようにする手順は次のとおりです。

  1. Compute Engine VM との接続を解除します。

    (vm)$ exit
    

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

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

    $ ctpu delete --name=retinanet-tutorial \
      --zone=europe-west4-a
    
  3. TPU の使用に対して不要な料金が発生しないように、ゾーンを指定している ctpu status を実行してインスタンスが割り当てられていないことを確認します。削除には数分かかることがあります。次のようなレスポンスは、割り当てられたインスタンスがないことを示します。

    $ ctpu status --name=retinanet-tutorial \
      --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:             --
    
  4. 次に示すように gsutil を実行します。bucket-name の部分は、このチュートリアルで作成した Cloud Storage バケット名に置き換えてください。

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

次のステップ

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

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

別のベースを使用する

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