Cloud TPU での TensorFlow の実行

このドキュメントでは、Cloud TPU で TensorFlow モデルを実行する方法について詳しく説明します。Cloud TPU で TensorFlow モデルをトレーニングする方法は、使用している Cloud TPU アーキテクチャによって異なります。詳細については、システム アーキテクチャをご覧ください。

概要

アプリケーションを実行すると、TensorFlow は計算グラフを生成し、gRPC を介して TPU ノードに送信します。TPU ノードは、計算グラフを JIT 方式でコンパイルし、プログラム バイナリを 1 つ以上の TPU デバイスに送信して実行します。多くの場合、モデルへの入力は Cloud Storage に保存されます。TPU ノードは、入力データをトレーニング用に 1 つ以上の TPU デバイスにストリーミングします。

次のブロック図は、ニューラル ネットワーク モデル、TPU Estimator と TensorFlow クライアント、TensorFlow サーバーと XLA コンパイラで構成された、Cloud TPU ソフトウェア アーキテクチャを示しています。

画像

Cloud TPU の TensorFlow API

Keras は、Python で記述された高レベルのディープ ラーニング API で、機械学習プラットフォームの TensorFlow 上で動作します。Cloud TPU を使用するニューラル ネットワーク モデルを作成する場合は、Keras を使用することをおすすめします。

TPU Estimator(レガシー)

TPU Estimator は、Estimator 上に構築された以前の高レベル API セットで、Cloud TPU 用のモデル構築を簡素化し、かつ、TPU のパフォーマンスを最大化します。TPU Estimator は引き続きサポートされますが、これではなく Keras を使用することをおすすめします。

TensorFlow クライアント

TPU Estimator により、プログラムは TensorFlow 演算に変換され、変換された演算は、TensorFlow クライアントによって計算グラフに変換されます。計算グラフは、TensorFlow クライアントにより TensorFlow サーバーに送信されます。

TensorFlow サーバー

TensorFlow サーバーは、Cloud TPU サーバー上で実行されます。サーバーで TensorFlow クライアントからの計算グラフを受信すると、サーバーでは次の操作が行われます。

  1. Cloud Storage から入力を読み込む
  2. グラフを Cloud TPU 上で実行できる部分と CPU 上で実行する必要がある部分に分割する
  3. Cloud TPU 上で実行するサブグラフに対応する XLA 演算を生成する
  4. XLA コンパイラを呼び出す

TPU ノードで TensorFlow モデルをトレーニングする方法の詳細については、Cloud TPU クイックスタートをご覧ください。

Cloud TPU VM

Cloud TPU VM にはユーザー VM がなく、モデルとコードは TPU ホストマシン上の VM で実行されます。TPU ホストには、直接 SSH で接続できます。TensorFlow クライアントとサービスは、同じ VM または異なる VM に作成できます。TPU ホストでは、任意のコードを実行して前処理などを行えます。

画像

TPU VM での TensorFlow モデルのトレーニングの詳細については、Cloud TPU VM Tensorflow クイックスタートをご覧ください。