TPU を使用したモデルのトレーニング

Tensor Processing Unit(TPU)は、機械学習ワークロードを高速化するために使用される、Google 独自開発の ASIC です。Cloud TPU を使用して、AI Platform でトレーニング ジョブを実行できます。AI Platform にはジョブ管理インターフェースが用意されているので、TPU を自分で管理する必要はありません。CPU や GPU でトレーニングを実行する場合と同じように AI Platform jobs API を使用できます。

高レベルの TensorFlow API は、Cloud TPU ハードウェア上でモデルを実行する場合に役立ちます。

GCP 環境の設定とテスト

スタートガイドのセットアップ セクションに従って、GCP 環境を構成します。

Cloud TPU にプロジェクトへのアクセスを許可する

GCP プロジェクトに関連付けられた Cloud TPU サービス アカウント名を承認するには、次の手順に従います。

  1. projects.getConfig を呼び出して、Cloud TPU サービス アカウント名を取得します。例:

    curl -H "Authorization: Bearer $(gcloud auth print-access-token)"  \
        https://ml.googleapis.com/v1/projects/<your-project-id>:getConfig
    
  2. API から返された tpuServiceAccount フィールドの値を保存します。

Cloud ML サービス エージェントの役割を使用して、Cloud TPU サービス アカウントをプロジェクト内のメンバーとして追加します。Google Cloud Platform Console または gcloud コマンドを使用して、次の手順に従います。

Console

  1. Google Cloud Platform Console にログインし、TPU を使用しているプロジェクトを選択します。
  2. [IAM と管理] > [IAM] を選択します。
  3. [追加] ボタンをクリックして、プロジェクトにメンバーを追加します。
  4. [メンバー] テキスト ボックスに TPU サービス アカウントを入力します。
  5. [役割] プルダウン リストをクリックします。
  6. Cloud ML サービス エージェントの役割を有効にします([サービス管理] > [Cloud ML サービス エージェント])。

gcloud

  1. プロジェクト ID と Cloud TPU サービス アカウントが格納されている環境変数を設定します。

    PROJECT_ID=your-project-id
    SVC_ACCOUNT=your-tpu-sa-123@your-tpu-sa.google.com.iam.gserviceaccount.com
    
  2. Cloud TPU サービス アカウントに ml.serviceAgent 役割を付与します。

    gcloud projects add-iam-policy-binding $PROJECT_ID \
        --member serviceAccount:$SVC_ACCOUNT --role roles/ml.serviceAgent
    

サービス アカウントへの役割の付与については、Cloud IAM のドキュメントをご覧ください。

ResNet-50 モデルのサンプルを実行

このセクションでは、gs://cloud-tpu-test-datasets/fake_imagenet にある架空のデータセットを使用して、リファレンスの TensorFlow ResNet-50 モデルをトレーニングします。このサンプルのジョブでは、マシン構成にあらかじめ定義された BASIC_TPU スケール階層を使用します。ガイドのこの後のセクションでは、カスタム構成を設定する方法について説明します。

AI Platform 上で以下のコマンドを実行して、コードを取得し、トレーニング ジョブを送信します。

  1. リファレンス モデル用のコードをダウンロードします。

    mkdir tpu-demos && cd tpu-demos
    wget https://github.com/tensorflow/tpu/archive/r1.14.tar.gz
    tar -xzvf r1.14.tar.gz && rm r1.14.tar.gz
    
  2. 解凍されたディレクトリ構造内の official ディレクトリに移動します。

    cd tpu-r1.14/models/official/
    
  3. ./resnet/resnet_main.py を編集し、サブモジュールのインポート時に明示的な相対インポートを使用するようにコードを変更します。たとえば、次のように記述されていたとします。

    import resnet_model
    

    これを以下のように変更します。

    from . import resnet_model
    

    上記のパターンを使用して、ファイル内の他のインポートをすべて変更してから、ファイルを保存します。

  4. サンプル内の他のファイルにあるサブモジュールのインポートを確認し、明示的な相対インポートも使用するように更新します。

  5. いくつかの環境変数を設定します。

    JOB_NAME=tpu_1
    STAGING_BUCKET=gs://my_bucket_for_staging
    REGION=us-central1
    DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
    OUTPUT_PATH=gs://my_bucket_for_model_output
    

    次のリージョンで現在 TPU にアクセスできるようになっています。

    • us-central1

  6. gcloud ai-platform jobs submit training コマンドを使用して、トレーニング ジョブを送信します。

    gcloud ai-platform jobs submit training $JOB_NAME \
            --staging-bucket $STAGING_BUCKET \
            --runtime-version 1.14 \
            --scale-tier BASIC_TPU \
            --module-name resnet.resnet_main \
            --package-path resnet/ \
            --region $REGION \
            -- \
            --data_dir=$DATA_DIR \
            --model_dir=$OUTPUT_PATH
    

Cloud TPU でのモデルのトレーニングの詳細

このガイドの前半では、ResNet-50 のサンプルコードの使い方について説明します。このセクションでは、AI Platform 上で Cloud TPU を使用してジョブの構成とモデルのトレーニングを行う方法について詳しく説明します。

TPU を提供するリージョンの指定

TPU が利用可能なリージョンでジョブを実行する必要があります。次のリージョンで現在 TPU にアクセスできるようになっています。

  • us-central1

モデル トレーニングやオンライン / バッチ予測などの AI Platform サービスで利用可能なリージョンの詳細については、リージョンのガイドをご覧ください。

TensorFlow と AI Platform のバージョニング

Cloud TPU でモデルをトレーニングする際に利用できる AI Platform ランタイムのバージョンには 1.12、1.13、1.14 があります。AI Platform ランタイムのバージョンと対応する TensorFlow のバージョンについてはこちらをご覧ください。

バージョニング ポリシーは Cloud TPU と同じです。トレーニングのジョブ リクエストでは、TPU で使用可能なランタイム バージョンを指定し、トレーニング コードで使用されている TensorFlow バージョンと一致することを確認してください。

TPU gRPC サーバーとの接続

TensorFlow プログラムでは、TPUClusterResolver を使用して、TPU VM 上で実行されている TPU gRPC サーバーに接続する必要があります。TPUClusterResolver は、Cloud TPU の IP アドレスとポートを返します。

次の例は、ResNet-50 のサンプルコードで TPUClusterResolver を使用する方法を示しています。

tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
    FLAGS.tpu,
    zone=FLAGS.tpu_zone,
    project=FLAGS.gcp_project)

config = tpu_config.RunConfig(
    cluster=tpu_cluster_resolver,
    model_dir=FLAGS.model_dir,
    save_checkpoints_steps=max(600, FLAGS.iterations_per_loop),
    tpu_config=tpu_config.TPUConfig(
        iterations_per_loop=FLAGS.iterations_per_loop,
        num_shards=FLAGS.num_cores,
        per_host_input_for_training=tpu_config.InputPipelineConfig.PER_HOST_V2))  # pylint: disable=line-too-long

TPU へのオペレーションの割り当て

マシン上の TPU を使用するには、TensorFlow TPUEstimator API を使用する必要があります。この API は上位の TensorFlow Estimator API を継承しています。

  • TPUEstimator は、TPU デバイスの実行に関するさまざまな詳細項目を処理します。たとえば、各コアの入力とモデルを複製したり、フックを実行するために定期的にホストに戻したりします。
  • 上位の TensorFlow API には、他にも多くの利点があります。特に、この API はモデルのチェックポイントを保存して復元するため、停止した時点で中断したトレーニング ジョブを再開できます。

Cloud TPU で利用可能な TensorFlow オペレーションのリストをご覧ください。

カスタム TPU マシンの構成

TPU トレーニング ジョブは、2 つの VM による構成で実行されます。一方の VM(マスター)が Python コードを実行します。マスターは、TPU ワーカーで実行されている TensorFlow サーバーを駆動します。

AI Platform で TPU を使用するには、次のいずれかの方法で TPU 対応マシンにアクセスするようにトレーニング ジョブを構成します。

  • BASIC_TPU スケール階層を使用する。この方法で TPU v2 アクセラレータにアクセスできます。
  • マスター VM で cloud_tpu ワーカーと従来のマシンタイプを使用する。この方法で TPU v2 アクセラレータにアクセスできます。
  • マスター VM で cloud_tpu ワーカーと Compute Engine マシンタイプを使用する。この方法で TPU v2 アクセラレータまたは TPU v3 アクセラレータにアクセスできます。TPU v3 アクセラレータはベータ版で利用できます。

基本的な TPU 対応マシン

スケール階層を BASIC_TPU に設定して、上記のサンプルを実行した場合と同様に、マスター VM と、8 つの TPU v2 コアを搭載した 1 つの TPU を使用している TPU VM を取得します。

従来のマシンタイプ構成の TPU ワーカー

また、マスター VM でさらに多くのコンピューティング リソースが必要な場合は、次のようにカスタムマシン構成を設定することもできます。

  • スケール階層を CUSTOM に設定する。
  • ジョブの要件に合った従来のマシンタイプを使用するようにマスター VM を構成する。
  • workerTypecloud_tpu に設定して、8 つの TPU v2 コアを搭載した 1 つの Cloud TPU を使用している TPU VM を取得します。
  • workerCount を 1 に設定する。
  • Cloud TPU を使用する場合は、パラメータ サーバーを指定しない。parameterServerCount がゼロより大きい場合、サービスはジョブ リクエストを拒否します。

次の例は、このタイプの構成を使用する config.yaml ファイルを示しています。

trainingInput:
  scaleTier: CUSTOM
  masterType: complex_model_m
  workerType: cloud_tpu
  workerCount: 1

Compute Engine マシンタイプ構成の TPU ワーカー

マスター VM の Compute Engine マシンタイプにカスタムマシン構成を設定し、TPU VM に接続された acceleratorConfig を設定することもできます。

このタイプの構成を使用して、8 つの TPU v2 コアを搭載した TPU ワーカー(acceleratorConfig なしの構成に類似)または 8 つの v3 コアを搭載した TPU ワーカー(ベータ版)を設定できます。TPU v2 アクセラレータと TPU v3 アクセラレータの違いをご確認ください。

Compute Engine マシンタイプを使用すると、次のようにマスター VM をより柔軟に構成できます。

  • スケール階層を CUSTOM に設定する。
  • ジョブ要件に合った Compute Engine マシンタイプを使用するようにマスター VM を構成する。
  • workerTypecloud_tpu に設定する。
  • acceleratorConfig フィールドを含む workerConfig を追加します。その acceleratorConfig 内で、typeTPU_V2 または TPU_V3 に、count8 に設定します。それ以外の数の TPU コアを接続することはできません。
  • workerCount を 1 に設定する。
  • Cloud TPU を使用する場合は、パラメータ サーバーを指定しない。parameterServerCount がゼロより大きい場合、サービスはジョブ リクエストを拒否します。

次の例は、このタイプの構成を使用する config.yaml ファイルを示しています。

TPU v2

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V2
      count: 8

TPU v3(ベータ版)

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V3
      count: 8

TPU ワーカーでのカスタム コンテナの使用

TPU をサポートする AI Platform ランタイム バージョンのいずれかを使用せずに、TPU ワーカーでカスタム コンテナを実行する場合は、トレーニング ジョブの送信時に追加の構成フィールドを指定する必要があります。tpuTfVersion を、コンテナで使用する TensorFlow のバージョンが含まれているランタイム バージョンに設定します。TPU でのトレーニングのために現在サポートされているランタイム バージョンを指定する必要があります。

カスタム コンテナを使用するジョブを構成しているので、AI Platform のトレーニングではトレーニング ジョブの実行時にこのランタイム バージョンの環境を使用しません。ただし、カスタム コンテナで使用する TensorFlow のバージョン用に TPU ワーカーを適切に準備できるように、AI Platform のトレーニングでこのフィールドが必要です。

次の例では、前のセクションと同様の TPU 構成の config.yaml ファイルを示しますが、この場合はマスター VM と TPU ワーカーが異なるカスタム ワーカーで動作します。

TPU v2

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/your-project-id/your-master-image-name:your-master-tag-name
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/your-project-id/your-worker-image-name:your-worker-tag-name
    tpuTfVersion: 1.14
    acceleratorConfig:
      type: TPU_V2
      count: 8

TPU v3(ベータ版)

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/your-project-id/your-master-image-name:your-master-tag-name
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/your-project-id/your-worker-image-name:your-worker-tag-name
    tpuTfVersion: 1.14
    acceleratorConfig:
      type: TPU_V3
      count: 8

トレーニング ジョブを送信するために gcloud beta ai-platform jobs submit training コマンドを使用する場合は、config.yaml ファイルにではなく --tpu-tf-version フラグを指定して tpuTfVersion API フィールドを指定できます。

詳細については、カスタム コンテナを使用した分散トレーニングをご覧ください。

次のステップ