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.12.tar.gz
    tar -xzvf r1.12.tar.gz && rm r1.12.tar.gz
    
  2. 解凍されたディレクトリ構造内の official ディレクトリに移動します。

    cd tpu-r1.12/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 ml-engine jobs submit training コマンドを使用してトレーニング ジョブを送信します。

    gcloud ml-engine jobs submit training $JOB_NAME \
            --staging-bucket $STAGING_BUCKET \
            --runtime-version 1.12 \
            --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 のサンプルコードの使い方について説明します。このセクションでは、Cloud TPU を使用した AI Platform でのジョブの構成とモデルのトレーニングについて詳しく説明します。

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

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

  • us-central1

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

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

Cloud TPU でモデルをトレーニングする際に利用できる AI Platform ランタイムのバージョンには 1.11 と 1.12 があります。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 スケール階層を使用する。
  • マスター VM で cloud_tpu ワーカーと AI Platform マシンタイプを使用する。
  • マスター VM で cloud_tpu ワーカーと Compute Engine マシンタイプを使用する。

基本的な TPU 対応マシン

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

AI Platform のマシンタイプ構成における TPU ワーカー

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

  • スケール階層を CUSTOM に設定する。
  • ジョブの要件に合った AI Platform マシンタイプを使用するようにマスター VM を構成する。
  • 1 つの Cloud TPU を含む TPU VM を取得するように、workerTypecloud_tpu に設定する。
  • 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 を設定することもできます。

現在、これは acceleratorConfig のない構成と同じ TPU リソース(8 つの TPU v2 コア)を提供します。Compute Engine マシンタイプを使用すると、マスター VM をより柔軟に構成できますが、このようなトレーニング ジョブの構成はベータ版のリリース段階です。

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

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

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

次のステップ

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

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

TensorFlow 用 AI Platform