Cloud TPU での TensorFlow の実行

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

概要

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

次のブロック図は、ニューラル ネットワーク モデル、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 のクイックスタートをご覧ください。

Cloud TPU VM

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

画像

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