Cloud TPU での TensorFlow の実行

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

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

Cloud TPU VM

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

画像

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

Cloud TPU ノード

Cloud TPU ノード アーキテクチャでは、TPU ホストは Compute Engine VM と TPU ノードで構成されます。Compute Engine VM は TensorFlow クライアントとモデルを実行します。VM は gRPC を介して TPU ノードと通信します。TPU ノードは TensorFlow サーバーと XLA コンパイラを実行します。

次のブロック図は、ニューラル ネットワーク モデル、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 のパフォーマンスが最大限に向上します。TensorFlow 1.x では TPU Estimator は引き続きサポートされていますが、代わりに Keras を使用することをおすすめします。

TensorFlow サーバー

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

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

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