Cloud TPU と GKE を使用した Mask R-CNN のトレーニング

このチュートリアルでは、Cloud TPU と GKE で Mask RCNN モデルをトレーニングする方法について説明します。

目標

  • データセットとモデルの出力を格納する Cloud Storage バケットを作成します。
  • Cloud TPU リソースを管理する GKE クラスタを作成します。
  • Kubernetes ジョブ仕様をダウンロードします。このジョブ仕様には、Cloud TPU 上の TensorFlow で Mask RCNN をトレーニングするために必要なリソースが規定されています。
  • GKE クラスタでジョブを実行し、モデルのトレーニングを開始します。
  • ログとモデル出力を確認します。

費用

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

  • Compute Engine
  • Cloud TPU
  • Cloud Storage

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

始める前に

  1. Google アカウントにログインします。

    Google アカウントをまだお持ちでない場合は、新しいアカウントを登録します。

  2. GCP プロジェクトを選択または作成します。

    [リソースの管理] ページに移動

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

    課金を有効にする方法について

  4. GKE で Cloud TPU を使用する際に、プロジェクトは Google Cloud Platform の課金対象となるコンポーネントを使用します。Cloud TPU の料金GKE の料金を確認して費用を見積もり、使い終わったら手順に沿ってリソースのクリーンアップを行います。

  5. GCP Console で、次の API を有効にします。

Cloud Storage バケットを作成する

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

  1. GCP Console の Cloud Storage ページに移動します。

    Cloud Storage ページに移動

  2. 次のオプションを指定して新しいバケットを作成します。

    • 任意の一意な名前
    • デフォルトのストレージ クラス: Regional
    • ロケーション: us-central1

Cloud TPU に Cloud Storage バケットへのアクセスを許可する

Cloud TPU に Cloud Storage オブジェクトに対する読み取り / 書き込みアクセス権を付与する必要があります。そのためには、Cloud TPU で使用するサービス アカウントに必要なアクセス権を付与します。ガイドに沿ってストレージ バケットに対するアクセス権を付与してください。

GKE でクラスタを作成する

gcloud コマンドライン ツールを使用して環境を設定し、Cloud TPU をサポートする GKE クラスタを作成する手順は以下のとおりです。

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

    Cloud Shell を開く

  2. GCP プロジェクトを指定します。

    $ gcloud config set project YOUR-CLOUD-PROJECT
    

    YOUR-CLOUD-PROJECT は、GCP プロジェクトの名前です。

  3. Cloud TPU リソースを使用する予定のゾーンを指定します。この例では、us-central1-b ゾーンを使用します。

    $ gcloud config set compute/zone us-central1-b
    

    Cloud TPU は次のゾーンで使用できます。

    米国

    TPU タイプ(v2) TPU v2 コア数 TPU メモリ合計 利用可能なゾーン
    v2-8 8 64 GiB us-central1-a
    us-central1-b
    us-central1-c
    us-central1-f TFRC のみ
    v2-32(ベータ版) 32 256 GiB us-central1-a
    v2-128(ベータ版) 128 1 TiB us-central1-a
    v2-256(ベータ版) 256 2 TiB us-central1-a
    v2-512(ベータ版) 512 4 TiB us-central1-a
    TPU タイプ(v3) TPU v3 コア数 TPU メモリ合計 利用可能なリージョン
    v3-8 8 128 GiB us-central1-a
    us-central1-b
    us-central1-f TFRC のみ

    ヨーロッパ

    TPU タイプ(v2) TPU v2 コア数 TPU メモリ合計 利用可能なゾーン
    v2-8 8 64 GiB europe-west4-a
    v2-32(ベータ版) 32 256 GiB europe-west4-a
    v2-128(ベータ版) 128 1 TiB europe-west4-a
    v2-256(ベータ版) 256 2 TiB europe-west4-a
    v2-512(ベータ版) 512 4 TiB europe-west4-a
    TPU タイプ(v3) TPU v3 コア数 TPU メモリ合計 利用可能なリージョン
    v3-8 8 128 GiB europe-west4-a
    v3-32(ベータ版) 32 512 GiB europe-west4-a
    v3-64(ベータ版) 64 1 TiB europe-west4-a
    v3-128(ベータ版) 128 2 TiB europe-west4-a
    v3-256(ベータ版) 256 4 TiB europe-west4-a
    v3-512(ベータ版) 512 8 TiB europe-west4-a
    v3-1024(ベータ版) 1024 16 TiB europe-west4-a
    v3-2048(ベータ版) 2048 32 TiB europe-west4-a

    アジア太平洋

    TPU タイプ(v2) TPU v2 コア数 TPU メモリ合計 利用可能なゾーン
    v2-8 8 64 GiB asia-east1-c
  4. gcloud container clusters コマンドを使用して、Cloud TPU をサポートするクラスタを GKE 上に作成します。GKE クラスタとそのノードプールは、Cloud TPU を使用できるゾーン内に作成する必要があります。該当するゾーンは、上記の環境変数に関するセクションに記載されています。次のコマンドは、tpu-models-cluster という名前のクラスタを作成します。

    $ gcloud container clusters create tpu-models-cluster \
    --cluster-version=1.13 \
    --scopes=cloud-platform \
    --enable-ip-alias \
    --enable-tpu
    

    上記のコマンドの各要素の意味は次のとおりです。

    • --cluster-version=1.13 は、クラスタで最新の Kubernetes 1.13 リリースが使用されるように指定します。バージョンは 1.13.4-gke.5 以降を使用する必要があります。
    • --scopes=cloud-platform を設定することで、クラスタ内のすべてのノードに、上記の YOUR-CLOUD-PROJECT で定義された GCP 内の Cloud Storage バケットに対するアクセス権が付与されます。クラスタとストレージ バケットが同じプロジェクト内に存在していなければなりません。デフォルトでは、ポッドはデプロイされているノードのスコープを継承します。このフラグにより、クラスタ内で実行しているすべてのポッドに cloud-platform スコープが設定されます。ポッド単位でアクセスを制限したい場合は、GKE のサービス アカウントで認証するためのガイドを参照してください。
    • --enable-ip-alias は、クラスタがエイリアス IP 範囲を使用することを示します。これは、GKE 上で Cloud TPU を使用する場合に必要です。
    • --enable-tpu は、クラスタが Cloud TPU をサポートしなければならないことを示します。

    コマンドの実行が完了すると、次のような確認メッセージが表示されます。

    kubeconfig entry generated for tpu-models-cluster.
    NAME                LOCATION       MASTER_VERSION   MASTER_IP      MACHINE_TYPE   NODE_VERSION   NUM_NODES  STATUS
    tpu-models-cluster  us-central1-b  1.13.6-gke.5     35.232.204.86  n1-standard-2  1.13.6-gke.5   3          RUNNING
    

トレーニング データを処理する

次に、ジョブを作成します。ジョブは、Google Kubernetes Engine において有限のタスクを表すコントローラ オブジェクトです。最初に作成する必要があるジョブは、Mask RCNN モデルのトレーニングに使用する COCO データセットをダウンロードして処理するジョブです。

  1. 次のように、シェル環境で download_and_preprocess_coco_k8s.yaml という名前のファイルを作成します。あるいは、GitHub からこのファイルをダウンロードすることもできます。

    # Download and preprocess the COCO dataset.
    #
    # Instructions:
    #   1. Follow the instructions on https://cloud.google.com/tpu/docs/kubernetes-engine-setup
    #      to create a Kubernetes Engine cluster. The Job must be running at least
    #      on a n1-standard-4 machine.
    #   2. Change the environment variable DATA_BUCKET below to the path of the
    #      Google Cloud Storage bucket where you want to store the training data.
    #   3. Run `kubectl create -f download_and_preprocess_coco_k8s.yaml`.
    
    apiVersion: batch/v1
    kind: Job
    metadata:
      name: download-and-preprocess-coco
    spec:
      template:
        spec:
          restartPolicy: Never
          containers:
          - name: download-and-preprocess-coco
            # The official TensorFlow 1.13 TPU model image built from https://github.com/tensorflow/tpu/blob/r1.13/tools/docker/Dockerfile.
            image: gcr.io/tensorflow/tpu-models:r1.13
            command:
              - /bin/bash
              - -c
              - >
                cd /tensorflow_tpu_models/tools/datasets &&
                bash download_and_preprocess_coco.sh /scratch-dir &&
                gsutil -m cp /scratch-dir/*.tfrecord ${DATA_BUCKET}/coco &&
                gsutil cp /scratch-dir/raw-data/annotations/*.json ${DATA_BUCKET}/coco
            env:
              # [REQUIRED] Must specify the Google Cloud Storage location where the
              # COCO dataset will be stored.
            - name: DATA_BUCKET
              value: "gs://<my-data-bucket>/data/coco"
            volumeMounts:
            - mountPath: /scratch-dir
              name: scratch-volume
          volumes:
          - name: scratch-volume
            persistentVolumeClaim:
              claimName: scratch-disk-coco
    ---
    apiVersion: v1
    kind: PersistentVolumeClaim
    metadata:
      name: scratch-disk-coco
    spec:
      accessModes:
        - ReadWriteOnce
      resources:
        requests:
          storage: 100Gi
    
  2. DATA_BUCKET という名前の環境変数を見つけます。この変数の value を Cloud Storage バケットのパスに変更します。

  3. 次のコマンドを実行して、クラスタにジョブを作成します。

    kubectl create -f download_and_preprocess_coco_k8s.yaml
    

ジョブが開始するまでに数分かかることがあります。次のコマンドを使用すると、クラスタで実行中のジョブのステータスを確認できます。

kubectl get pods -w

このコマンドは、クラスタ内で実行されているすべてのジョブの情報を返します。例:

NAME                                 READY     STATUS      RESTARTS   AGE
download-and-preprocess-coco-xmx9s   0/1       Pending     0          10s

-w フラグを指定すると、ポッドのステータスの変化がモニタリングされます。数分後、ジョブのステータスが Running に更新されます。

ジョブに問題が発生した場合は、ジョブを削除して再試行できます。その場合、kubectl create コマンドを再度実行する前に、kubectl delete -f download_and_preprocess_coco_k8s.yaml コマンドを実行します。

Mask RCNN モデルを実行する

Cloud TPU と GKE を使用して、Mask RCNN モデルを実行する準備がすべて整いました。

  1. シェル環境で、次のように mask_rcnn_k8s.yaml という名前のファイルを作成します。あるいは、GitHub からこのファイルをダウンロードすることもできます。

    # Train Mask-RCNN with Coco dataset using Cloud TPU and Google Kubernetes Engine.
    #
    # [Training Data]
    #   Download and preprocess the COCO dataset using https://github.com/tensorflow/tpu/blob/r1.13/tools/datasets/download_and_preprocess_coco_k8s.yaml
    #   if you don't already have the data.
    #
    # [Instructions]
    #   1. Follow the instructions on https://cloud.google.com/tpu/docs/kubernetes-engine-setup
    #      to create a Kubernetes Engine cluster.
    #   2. Change the environment variable MODEL_BUCKET in the Job spec to the
    #      Google Cloud Storage location where you want to store the output model.
    #   3. Run `kubectl create -f mask_rcnn_k8s.yaml`.
    
    apiVersion: batch/v1
    kind: Job
    metadata:
      name: mask-rcnn-gke-tpu
    spec:
      template:
        metadata:
          annotations:
            # The Cloud TPUs that will be created for this Job must support
            # TensorFlow 1.13. This version MUST match the TensorFlow version that
            # your model is built on.
            tf-version.cloud-tpus.google.com: "1.13"
        spec:
          restartPolicy: Never
          containers:
          - name: mask-rcnn-gke-tpu
            # The official TensorFlow 1.13 TPU model image built from https://github.com/tensorflow/tpu/blob/r1.13/tools/docker/Dockerfile.
            image: gcr.io/tensorflow/tpu-models:r1.13
            command:
            - /bin/sh
            - -c
            - >
                DEBIAN_FRONTEND=noninteractive apt-get update &&
                DEBIAN_FRONTEND=noninteractive apt-get install -y python-dev python-tk libsm6 libxrender1 libxrender-dev libgtk2.0-dev libxext6 libglib2.0 &&
                pip install Cython matplotlib opencv-python-headless &&
                pip install 'git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI' &&
                python /tensorflow_tpu_models/models/official/mask_rcnn/mask_rcnn_main.py
                --model_dir=${MODEL_BUCKET}
                --params_override=iterations_per_loop=500,resnet_checkpoint=${RESNET_CHECKPOINT},resnet_depth=50,precision=bfloat16,train_batch_size=64,eval_batch_size=8,training_file_pattern=${DATA_BUCKET}/train-*,validation_file_pattern=${DATA_BUCKET}/val-*,val_json_file=${DATA_BUCKET}/instances_val2017.json,total_steps=22500
            env:
              # The Google Cloud Storage location to store dataset.
            - name: DATA_BUCKET
              value: "gs://<my-data-bucket>"
            - name: MODEL_BUCKET
              value: "gs://<my-model-bucket>/mask_rcnn"
            - name: RESNET_CHECKPOINT
              value: "gs://cloud-tpu-artifacts/resnet/resnet-nhwc-2018-02-07/model.ckpt-112603"
    
              # Point PYTHONPATH to the top level models folder
            - name: PYTHONPATH
              value: "/tensorflow_tpu_models/models"
            resources:
              limits:
                # Request a single v3-8 Cloud TPU device to train the model.
                # A single v3-8 Cloud TPU device consists of 4 chips, each of which
                # has 2 cores, so there are 8 cores in total.
                cloud-tpus.google.com/v3: 8
    
  2. DATA_BUCKET という名前の環境変数を見つけます。この変数の value を Cloud Storage バケットのパスに変更します。

  3. MODEL_BUCKET という名前の環境変数を見つけます。この変数の value を Cloud Storage バケットのパスに変更します。ここは、スクリプトがトレーニングの出力を保存する場所です。

  4. 次のコマンドを実行して、クラスタにジョブを作成します。

    kubectl create -f mask_rcnn_k8s.yaml
    

    このコマンドを実行すると、次の確認メッセージが表示されます。

    job "mask-rcnn-gke-tpu" created

トレーニング データを処理するで説明したように、ジョブの開始には数分かかることがあります。次のコマンドを使用して、クラスタで実行中のジョブのステータスを確認します。

kubectl get pods -w

トレーニングのステータスを確認する

kubectl コマンドライン ユーティリティを使用すると、トレーニング ステータスを追跡できます。

  1. 次のコマンドを実行して、ジョブのステータスを取得します。

    kubectl get pods
    

    トレーニング プロセスの実行中、ポッドのステータスは RUNNING になります。

  2. トレーニング プロセスの詳細情報を取得するには、次のコマンドを実行します。

    $ kubectl logs job name

    job name は完全なジョブ名です。例: mask-rcnn-gke-tpu-abcd

    GCP Console の GKE ワークロード ダッシュボードで出力を確認することもできます。

    ログに最初のエントリが表示されるまでには時間がかかるのでご注意ください。次のような出力が表示されます。

    I0622 18:14:31.617954 140178400511808 tf_logging.py:116] Calling model_fn.
    I0622 18:14:40.449557 140178400511808 tf_logging.py:116] Create CheckpointSaverHook.
    I0622 18:14:40.697138 140178400511808 tf_logging.py:116] Done calling model_fn.
    I0622 18:14:44.004508 140178400511808 tf_logging.py:116] TPU job name worker
    I0622 18:14:45.254548 140178400511808 tf_logging.py:116] Graph was finalized.
    I0622 18:14:48.346483 140178400511808 tf_logging.py:116] Running local_init_op.
    I0622 18:14:48.506665 140178400511808 tf_logging.py:116] Done running local_init_op.
    I0622 18:14:49.135080 140178400511808 tf_logging.py:116] Init TPU system
    I0622 18:15:00.188153 140178400511808 tf_logging.py:116] Start infeed thread controller
    I0622 18:15:00.188635 140177578452736 tf_logging.py:116] Starting infeed thread controller.
    I0622 18:15:00.188838 140178400511808 tf_logging.py:116] Start outfeed thread controller
    I0622 18:15:00.189151 140177570060032 tf_logging.py:116] Starting outfeed thread controller.
    I0622 18:15:07.316534 140178400511808 tf_logging.py:116] Enqueue next (100) batch(es) of data to infeed.
    I0622 18:15:07.316904 140178400511808 tf_logging.py:116] Dequeue next (100) batch(es) of data from outfeed.
    I0622 18:16:13.881397 140178400511808 tf_logging.py:116] Saving checkpoints for 100 into gs://<my-model-bucket>/mask_rcnn/model.ckpt.
    I0622 18:16:21.147114 140178400511808 tf_logging.py:116] loss = 1.589756, step = 0
    I0622 18:16:21.148168 140178400511808 tf_logging.py:116] loss = 1.589756, step = 0
    I0622 18:16:21.150870 140178400511808 tf_logging.py:116] Enqueue next (100) batch(es) of data to infeed.
    I0622 18:16:21.151168 140178400511808 tf_logging.py:116] Dequeue next (100) batch(es) of data from outfeed.
    I0622 18:17:00.739207 140178400511808 tf_logging.py:116] Enqueue next (100) batch(es) of data to infeed.
    I0622 18:17:00.739809 140178400511808 tf_logging.py:116] Dequeue next (100) batch(es) of data from outfeed.
    I0622 18:17:36.598773 140178400511808 tf_logging.py:116] global_step/sec: 2.65061
    I0622 18:17:37.040504 140178400511808 tf_logging.py:116] examples/sec: 2698.56
    I0622 18:17:37.041333 140178400511808 tf_logging.py:116] loss = 2.63023, step = 200 (75.893 sec)
    
  3. gs://<my-model-bucket>/mask-rcnn/model.ckpt でトレーニング済みのモデルを確認します。

    gsutil ls -r gs://my-model-bucket/mask-rcnn/model.ckpt

クリーンアップ

GKE 上で Cloud TPU を使い終わったら、リソースをクリーンアップして、Google Cloud Platform アカウントに余分に課金されないようにしてください。

このセッションでプロジェクトとゾーンを設定していない場合は、この時点で設定します。このガイドの上述の手順をご覧ください。続いて、次のクリーンアップ手順を実行します。

  1. tpu-models-cluster コマンドを実行して GKE クラスタを削除します。YOUR-PROJECT は、該当する GCP プロジェクト名で置き換えてください。

    $ gcloud container clusters delete tpu-models-cluster --project=YOUR-PROJECT
    
  2. データの検証が終わったら、gsutil コマンドを使用して、このチュートリアルで作成した Cloud Storage バケットを削除します。YOUR-BUCKET は、該当する Cloud Storage バケットの名前で置き換えてください。

    $ gsutil rm -r gs://YOUR-BUCKET
    

    無料の保存容量の制限やその他の料金に関する情報については、Cloud Storage の料金ガイドをご覧ください。

次のステップ

  • TensorBoard の TPU ツールを確認する。
  • 次のジョブ仕様のいずれかを使用して、他のモデルとデータセット取得ジョブを実行する。

    • GKE に COCO データセットをダウンロードして前処理します。
    • GKE に ImageNet をダウンロードして前処理します。
    • Cloud TPU と GKE を使用して AmoebaNet-D をトレーニングします。
    • Cloud TPU と GKE を使用して Inception v3 をトレーニングします。
    • Cloud TPU と GKE を使用して RetinaNet をトレーニングします。
  • 他の TPU サンプルで試してみる。

このページは役立ちましたか?評価をお願いいたします。

フィードバックを送信...