トラブルシューティングとよくある質問

このガイドでは、Cloud TPU で独自の TensorFlow モデルを実行するユーザーに、トラブルシューティングのヘルプを提供します。Cloud TPU を使い始めるためのより一般的なガイドについては、クイックスタートまたは MNIST チュートリアルをご覧ください。

概要

Cloud TPU で TensorFlow モデルを実行するためにおすすめする戦略は、TPUEstimator API を使用することです。TensorFlow の Estimator API をすでに使用している場合、TPUEstimator への切り替えには、通常は数行のコードの変更のみが必要です。TPUEstimator へのデータのロードに推奨される方法は、Dataset API の使用です。Dataset と TPUEstimator の使用方法の実例については、ResNet チュートリアルをご覧ください。

モデルを TPUEstimator に変換したら、フラグ use_tpu=False で動作することを確認することをおすすめします。このフラグにより、TensorFlow は通常の Estimator API にフォールバックし、TPU に関連するコードは使用されなくなります。したがって、use_tpu=False でモデルを実行する際に発生する問題は TPU には関係なく、このガイドの範囲外です。TensorFlow に関する一般的なヘルプについては、TensorFlow のドキュメントをご覧ください。

理想的には、TPUEstimator と use_tpu=False を使用してモデルを正常に実行できると、それを TPU で実行するには、use_tpu=True を設定し、master が TPU サーバー URL を指すようにするだけです（通常はクラスタリゾルバを使用）。ただし、TensorFlow モデルは非常に複雑な場合があり、TPU ではまったく新しい実行エンジンが使用されるため、TPU 固有の問題が発生する可能性があります。問題は次の 5 つの大きなカテゴリに分類されます（このガイドの関連するセクションへのリンクがあります）。

トレーニングスクリプトが TPU サーバーにまったく接続できない。
モデルを実行しようとすると TPU からエラーが返される。
モデルが TPU メモリに収まらない。
モデルを TPU で実行できるが、トレーニング速度が予期したほど速くない。
モデルを TPU で実行できるが、TPU トレーニング済みモデルの精度が CPU / GPU トレーニングのベースラインよりも悪い。

さらに、このガイドには TPU で利用可能な一般的な機能に関するよくある質問が含まれています。

特定のタイプのニューラルネットワークを TPU に移植するための特別なヘルプについては、Cloud TPU チュートリアルをご覧ください。

TPU サーバーへの接続の問題

TPU でモデルを実行するときは、リモート TPU サーバーの URL を RunConfig の master パラメータに渡す必要があります。内部では、TensorFlow によって、このサーバーとのリモート tf.Session が作成されます。このセクションでは、TPU サーバーに接続するときに TensorFlow がハングしたりエラーを出力したりする状況のトラブルシューティングについて説明します。大規模なモデルの場合は TPU グラフのコンパイルステップに時間がかかることがあるため、スクリプトがハングしていると判断する前に、スクリプトを少なくとも 5 分間実行しください。

最初のステップは、問題がサーバー自体にあるのか、TensorFlow トレーニングパイプラインにあるのかを確認することです。これを行うには、TPU サーバー URL を使用して MNIST チュートリアルを実行し、これが正しく動作することを確認します。MNIST チュートリアルとの接続の問題がまだある場合は、TPU サーバーの問題であることが確認されます。この場合、次の手順に従います。

次のコマンドを実行して、使用可能な TPU を一覧表示します。
```
(vm)$ gcloud compute tpus list
```
MNIST チュートリアルに示されているように、zone と project も設定する必要がある場合があります。これにより、次のような出力が表示されます。
```
NAME       ZONE           ACCELERATOR_TYPE  NETWORK_ENDPOINT   NETWORK  RANGE          STATUS
demo-tpu   us-central1-b  v2-8              10.240.1.2:8470    default  10.240.1.0/29  READY
```
正しい値を --tpu_name に渡していること（上記の例では demo-tpu）、この TPU が READY と表示されていることを確認します。また、以下を使用して、zone と project が設定されていることも確認します。
```
(vm)$ gcloud config set project your-project-name
```
```
(vm)$ gcloud config set compute/zone us-central1-b
```
TPU が READY と表示されない場合や、接続にまだ問題がある場合は、gcloud compute tpus reset $TPU_NAME を使用してサーバーを手動で再起動します。上記の例では、$TPU_NAME は demo-tpu です。この処理には数分かかることがあります。
上記の ... tpus list コマンドを再実行し、TPU が READY 状態になるのを待ちます。この処理には数分かかることがあります。
MNIST チュートリアルを再度実行してみます。
まだ MNIST チュートリアルの実行に問題がある場合は、TPU サポートにお問い合わせください。

MNIST の例が正しく実行されるがモデルがまだハングする場合、問題はトレーニングパイプラインにある可能性があります。まず、モデルが TPUEstimator API を使用していることを確認します。この API は、複雑な処理パイプラインを処理するだけでなく、use_tpu フラグを使用して TPU と TPU 以外の実行を簡単に切り替えられるようにしているためです。TPUEstimator の使用方法の例については、TPU チュートリアルをご覧ください。TPUEstimator API を使用しているモデルが、use_tpu=False が設定されているときに正しく実行されることを確認してください。use_tpu=False が設定されているときにモデルが正しく実行されない場合、問題は TPU とは関係ありません。

一般的なエラーのデバッグ

ローカルファイルシステムを使用できない

エラーメッセージ

InvalidArgumentError: Unimplemented: File system scheme '[local]' not implemented

詳細

すべての入力ファイルとモデルディレクトリは Cloud Storage バケットパス（gs://bucket-name/...）を使用する必要があり、このバケットは TPU サーバーからアクセス可能である必要があります。すべてのデータ処理とモデルチェックポインティングは、ローカルマシンではなく TPU サーバー上で実行されることに注意してください。TPU で使用するために Cloud Storage を適切に構成する方法については、Cloud Storage バケットへの接続ガイドを参照してください。

サポートされないデータ型

エラーメッセージ

TypeError: DataType is not a supported TPU infeed type.

詳細

現在、TPU では tf.float32、tf.int32、tf.bfloat16、tf.bool データ型のみがサポートされています。その他の一般的なデータ型（tf.uint8、tf.string、tf.int64 など）は、データ前処理中に（つまり、TPUEstimator の input_fn で）サポートされているデータ型のいずれかに変換する必要があります。別の例については、MNIST チュートリアルをご覧ください。一例として、MNIST の次のコードスニペットでは、tf.uint8 バイトシーケンスとして保存された image テンソルが tf.float32 テンソルに変換されます。

image = tf.decode_raw(image, tf.uint8)
image = tf.cast(image, tf.float32)
image = tf.reshape(image, [784])

次のスニペットでは、tf.int64 として保存された label テンソルが tf.int32 テンソルに変換されます。

label = tf.cast(label, tf.int32)

動的形状がサポートされない

エラーメッセージ

ValueError: shape [Shape] must have a fixed size for dimension d that is known at graph construction time.

詳細

TPU でモデルを実行するために、TensorFlow はXLA フレームワークを使用してモデルをコンパイルします。このコンパイルステップにより、トレーニング速度とメモリ使用量は大幅に向上しますが、グラフ内のすべてのテンソルの形状（ディメンションサイズ）は静的でなければなりません。つまり、グラフのコンパイル時にその値が既知である必要があります。コンパイル時に形状を判別できない場合、上記のようなエラーで TPU コンパイルは失敗します。

動的形状が返される一般的な演算の 1 つは、dataset.batch(batch_size) です。これは、ストリームに残っているサンプル数がバッチサイズよりも小さい可能性があるためです。したがって、TPU でのトレーニングでは tf.contrib.data.batch_and_drop_remainder(batch_size) を使用します。この場合、すべてのバッチが batch_size という静的形状になるように、最後のいくつかのサンプルがファイルから削除される可能性があります。次に例を示します。

dataset = ...
dataset = dataset.apply(tf.contrib.data.batch_and_drop_remainder(batch_size))

利用できない TensorFlow 演算

エラーメッセージ

NotFoundError: No registered 'OpName' OpKernel for XLA_TPU_JIT devices compatible with node

詳細

TPU で現在は利用できない TensorFlow 演算がモデルで使用されています。

TPU で利用可能な演算のリストと、今後のサポート計画や回避策の提案については、利用可能な TensorFlow 演算のガイドをご覧ください。

メモリ不足のエラーメッセージ

エラーメッセージ

ResourceExhaustedError: Ran out of memory in memory space hbm; used: YYY; limit: 7.48G.

詳細

各 Cloud TPU は 8 個の TPU コアで構成され、それぞれに 8 GB の RAM（HBM、高帯域メモリ）が搭載されています。このメモリは、勾配計算に必要な中間結果のテンソルだけでなく、重み（変数）テンソルを保存するために使用されます。モデルが大きすぎて TPU RAM に収まらない場合、初期化は失敗し、上記のエラーメッセージが表示されます。詳細については、メモリ使用量の削減をご覧ください。

`CrossShardOptimizer` を使用していない

エラーメッセージ

ValueError: CrossShardOptimizer must be used for model training on TPUs.

詳細

TensorFlow Python API を使用してモデルを定義する場合、ユーザーが記述するコードの大部分は、TPU に特化させる必要はありません。最も重要な例外はオプティマイザーであり、これは次に示すように tf.contrib.tpu.CrossShardOptimizer() でラップする必要があります。

optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate)
if FLAGS.use_tpu:
  optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)
train_op=optimizer.minimize(loss, tf.train.get_global_step())

各 Cloud TPU は 8 個の TPU コアで構成され、それらは独立したプロセッシングユニットです。各トレーニングステップ（すなわち、重み更新）では、各 TPU コアによりデータの独立したミニバッチ上でフォワードパスおよび勾配計算が実行され、その後、すべてのコアによって相互に勾配が交換されます。ほとんどの場合、これは数学的には 1 つの大きなバッチの勾配を計算することと同等ですが、データシャーディングについてで説明されている注意点があります。

CrossShardOptimizer は、この勾配交換を行う演算です。デフォルトでは、CrossShardOptimizer によってコア間の平均損失の勾配が計算されますが、合計損失を計算するように構成できます。そのためには、reduction=losses.Reduction.SUM を渡します。

TPU サーバーに接続できない

エラーメッセージ

An error was raised while a session was being created. This may be due to a preemption of a connected worker or parameter server. A new session is created.

詳細

このエラーは、master に渡された TPU サーバー URL に TensorFlow が接続できない場合に表示されます。詳細については、TPU サーバーへの接続の問題をご覧ください。

トレーニング中のエラー

TPU でモデルを正常に実行できない場合、これに関連するエラーは初期化中にキャッチされるように設計されています。したがって、トレーニング中にモデルが失敗することはまれです。これが発生した場合、最も考えられる原因はデータ前処理機能での問題です。たとえば、Dataset API を使用する場合、通常は dataset = dataset.repeat() を呼び出す必要があります。そうしないと、データを 1 回処理した後にトレーニングは失敗します。tf.while_loop() のような動的実行演算も、入力データに依存する方法でのみ、失敗することがあります。ハードウェアやネットワークの見かけ上の障害という、まれな可能性もあります。

実行の停止を起こす問題

TPU の実行中に TensorFlow でエラーが発生した場合、スクリプトはシェルを終了するのではなくハングしているように見えることがあります。これが発生した場合は、キーボードの CTRL+\ を押して SIGQUIT をトリガーします。これにより、Python はすぐに終了します。

同様に、TPU の実行中に CTRL+C を押すと、TensorFlow はすぐにはシャットダウンされず、現在の反復ループの終了を待って正常に終了します。CTRL+\ を押すと、Python はすぐに終了します。

この方法で終了した後に TPU サーバーに再接続できない場合は、gcloud compute tpus reset $TPU_SERVER_NAME コマンドを使用して TPU サーバーを手動でリセットします。$TPU_SERVER_NAME は gcloud compute tpus list コマンドの最初の列から取得します。

メモリ使用量の削減

TPU でモデルを実行しているときにメモリ不足エラーが発生した場合は、モデルのメモリ使用量を減らすための手順を取る必要があります。このセクションでは、メモリの問題の根本原因について説明し、それらを修正するためのガイドラインを示します。

モデルの重みの数が多い

メモリの問題の考えられる原因

float32 モデルの重みには、それぞれ 4 バイト必要です。これらの重みは、各 TPU コアで複製されます。したがって、数億の重みを持つモデルは、大きすぎて TPU に収まらない可能性があります。

メモリ使用量を減らす方法

特定のオプティマイザーでは、更新統計を保存するために重みごとに追加のメモリが必要です。特に、AdamOptimizer と AdadeltaOptimizer はどちらも、重みごとに 8 バイトが追加で必要です。AdagradOptimizer と MomentumOptimizer は、重みごとに 4 バイトが追加で必要です。標準的な GradientDescentOptimizer では、追加のストレージは必要ありませんが、最終的なモデルの精度に関して他のオプティマイザーよりも劣る場合があります。試験運用版の AdafactorOptimizer は、追加のメモリをほとんど必要とせず、Transformer モデルのトレーニング時にベースラインの Adam オプティマイザーと同様に実行されます。
重みの大部分が単語埋め込みである場合、WordPiece などの手法により、語彙のサイズが大幅に縮小され、さまざまなタスクで精度が向上することが示されています。
TensorFlow の今後のリリースには、16 ビット浮動小数点の重みと勾配の試験運用版サポートがあり、メモリ要件は半減する予定です。

テンソルパディングが過剰である

メモリの問題の考えられる原因

TPU メモリ内のテンソルはパディングされます。つまり、計算をより効率的に実行するために、メモリに保存されたテンソルのサイズは TPU によって切り上げられます。このパディングはハードウェアレベルで透過的に行われ、結果には影響しません。ただし、場合によっては、パディングによってメモリ使用量と実行時間が大幅に増加する可能性があります。

メモリ使用量を減らす方法

TPU ソフトウェアは、計算効率を最大にし、パディングを最小限に抑えるように、メモリ内にテンソルをレイアウトしようとします。ただし、このメモリレイアウトプロセスは複雑であり、最適な結果を得るにはモデルが以下の経験則に従っている必要があります。メモリのオーバーヘッドを最小限に抑え、計算効率を最大にするには、次のいずれかに該当する必要があります。

バッチサイズの合計は 64 の倍数（TPU コアあたり 8）である必要があり、フィーチャディメンションは 128 の倍数である必要があります。

または
バッチサイズの合計は 1024 の倍数（TPU コアあたり 128）である必要があり、フィーチャディメンションは 8 の倍数である必要があります。

バッチサイズ 1024 および 128 の倍数であるフィーチャディメンションを使用すると、最大の効率性が得られます。ただし、これはすべてのモデルで可能というわけではない場合があります。明確には、「フィーチャディメンション」は、全結合層の隠れサイズまたは畳み込みにおける出力チャネルの数を指します。すべての層がこのルールに準拠できるわけではありません（特にネットワークの最初と最後の層）。これは問題なく、ほとんどのモデルでは一定量のパディングが必要になると想定されます。

バッチサイズが大きすぎる

メモリの問題の考えられる原因

CPU、GPU、TPU でニューラルネットワークをトレーニングする場合、メモリの使用は 2 つの場所から発生します。

重み、重みの勾配、オプティマイザー固有の統計（慣性など）の保存。メモリの使用は、バッチサイズではなく、モデル内の重みの数に正比例します。
バックワードパスを計算するために必要なフォワードパスからの中間アクティベーションの保存。メモリの使用は、バッチサイズ、層サイズ、層の数に正比例します。

したがって、モデルに必要なメモリはバッチサイズに大きく依存します。

メモリ使用量を減らす方法

メモリ内に収まるまでバッチサイズを徐々に減らし、バッチサイズの合計が 64 の倍数になるようにします（コアあたりのバッチサイズは 8 の倍数である必要があります）。バッチサイズが大きいほど TPU でより効率的であることに注意してください。バッチサイズの合計 1024（コアあたり 128）から始めるのが、一般的には適切です。

モデルが大きすぎる

メモリの問題の考えられる原因

モデルに必要なメモリは、グラフ内の演算子（つまり、ネットワーク内の層）の数に大きく依存します。このストレージ要件は、重みの数とは別です。たとえば、tf.nn.conv2d() のような演算子の勾配を計算すると、重みを保存するために使用されるメモリに加えて、メモリの使用が増加する場合があります。

TPU エンジンは、モデルをメモリ内に収めるために特定の演算子を戦略的に再計算しようとしますが（再実体化と呼ばれ、勾配チェックポインティングと似ています）、これが常に可能とは限りません。

メモリ使用量を減らす方法

小さなバッチサイズ（64 など）でもモデルを TPU で実行できない場合は、層の数または層サイズを減らしてみてください。TensorFlow の今後のリリースでは、TPU での「モデルの並列処理」がサポートされる予定です。これにより、モデルの異なる部分を異なる TPU コアで実行することで、非常に大きいモデルを Cloud TPU で実行できるようになります。

トレーニング速度の向上

モデルを TPU で正常に実行できるが、トレーニング速度が予期したほど速くない場合があります。このセクションでは、速度を向上させる可能性があるいくつかの方法の概略を示します。

使用していない TPU コアがある

パフォーマンスの問題の説明

各 Cloud TPU には 8 個の別々の TPU コアが含まれており、それらは独立したプロセッシングユニットとして動作します。8 個のコアすべてを使用していない限り、TPU は完全には利用されていません。

モデルに影響しているかどうかを確認する方法

モデルで TPUConfig の num_shards パラメータを明示的に 8 に指定していない場合、TensorFlow はすべての TPU コアを利用するわけではありません。

緩和する方法

本格的な TPU トレーニングを行う場合は、常に num_shards を 8 に設定します。例については、MNIST チュートリアルまたは以下のコードスニペットをご覧ください。

tf.flags.DEFINE_integer("num_shards", 8, "Number of shards (TPU chips).")

run_config = tf.contrib.tpu.RunConfig(
  tpu_config=tf.contrib.tpu.TPUConfig(num_shards=FLAGS.num_shards, ...

estimator = tf.contrib.tpu.TPUEstimator(
    config=run_config, ...

num_shards を 1 に設定すると、モデルの精度の差異をデバッグするために役立つ場合があるため、デフォルト値 8 でコマンドラインフラグを使用することをおすすめします。

ループあたりの反復回数が少なすぎる

パフォーマンスの問題の説明

TPUConfig の iterations_per_loop パラメータによって、1 回の「トレーニングループ」で TPU に送信されるデータのバッチ数が制御されます。各トレーニングループでは、ローカルマシンと TPU サーバーの間で大量の通信が必要となるため、iterations_per_loop が小さすぎるとトレーニングが大幅に遅くなる可能性があります。

モデルに影響しているかどうかを確認する方法

ロギングメッセージ Enqueue next (X) batch(es) of data to infeed が頻繁に（たとえば 3 秒ごとに）表示される場合、トレーニングでトレーニングループから重大なオーバーヘッドが発生している可能性があります。

緩和する方法

iterations_per_loop を大きい値に設定します。MNIST チュートリアルでは、これは --iterations フラグによって制御されます。Enqueue next (X) batch(es) of data to infeed メッセージが 1 分間に数回以上表示されなくなれば、現在の値で十分です。iterations_per_loop を非常に大きい値に設定できますが、唯一の欠点は、ロギングメッセージとチェックポインティングはループの最後にのみ発生することです。

入力処理のボトルネック

パフォーマンスの問題の説明

TPU が特定のデータチャンクでトレーニングしている間に、入力処理関数によって CPU で次のデータチャンクが準備されます。したがって、入力関数がモデル関数よりも時間がかからない場合、入力処理のコストは事実上ゼロになります。しかし、入力関数がモデル関数よりも時間がかかると、ボトルネックが発生します。

モデルに影響しているかどうかを確認する方法

TensorBoard で入力パイプライン分析を表示するには、Cloud TPU ツール: Input Pipeline Analyzer の指示に従います。

入力パイプライン分析ページには、モデルで入力処理がボトルネックになっているかどうかを示す明確な概要が表示されます。同じページに演算ごとの実行時間も表示され、問題のある演算を特定できます。

緩和する方法

Dataset API を使用してデータをロードするときに、いくつかの緩和策があります。

データを tf.train.Example 構造のコレクションとして TFRecord ファイルに保存し、それを TFRecordDataset を使用してロードします。例については、Dataset API チュートリアルまたは ResNet チュートリアルをご覧ください。
dataset.cache() や dataset.prefetch() を使用して入力データをバッファします。これにより、ファイルアクセスが散発的に遅くなってボトルネックが発生するのを防ぎます。
dataset.map() 関数の num_parallel_calls パラメータを指定して、マルチスレッドの map() 演算を有効にします。
すべてのトレーニングのすべてのエポックでコストをかけるのではなく、コストのかかるデータ前処理を 1 回限りのコストとしてオフラインで実行します。

すべての入力処理はローカルマシン上ではなく TPU サーバー上にある CPU で実行されるため、ローカルマシンの速度は要因ではありません。

非行列乗算演算が多すぎる

パフォーマンスの問題の説明

Cloud TPU は、行列乗算と畳み込みを非常に高速に実行できます。他の TensorFlow 演算の大部分は TPU に効率的に実装されていますが、これらは他のハードウェアに対する TPU の主要な強みではありません。したがって、TPU を十分に利用するには、モデルで行列乗算または畳み込みが多くを占めている必要があります。

モデルに影響しているかどうかを確認する方法

Cloud TPU ツール: Op Profile ガイドでは、モデルのパフォーマンスプロファイルを演算タイプ別に生成する方法について説明されています。一般に、最新のニューラルネットワークアーキテクチャの大半では、行列乗算と畳み込みが多くを占めています。

緩和する方法

モデル内の行列乗算の不足が、主に他のハードウェアでのトレーニング速度の問題が起因となった場合は、速度のパフォーマンス向上のために、それらのモデルを TPU で再ベンチマークすることをおすすめします。モデル内の行列乗算の不足がモデルの基本的な特性である場合、TPU は最適なハードウェアの選択肢ではない可能性があります。

テンソルパディングが過剰である

パフォーマンスの問題の説明

TPU では、その計算ユニットを効率的に使用できるように、メモリ内のテンソルはパディングされます。パディングによって、メモリとメモリ帯域幅の両方の使用量が増加する可能性があります。テンソルパディングの問題を理解して解決するには、テンソルパディングをご覧ください。

バッチサイズが小さすぎる

パフォーマンスの問題の説明

一般的に、使用するバッチサイズが大きいほど、サンプル数 / 秒に関して TPU でのトレーニング速度は向上します。

モデルに影響しているかどうかを確認する方法

モデルのバッチサイズは、常に少なくとも 64（TPU コアあたり 8）である必要があります。これは、TPU ではテンソルは常にこのサイズにパディングされるためです。TPU でのトレーニングの理想的なバッチサイズは、1024（TPU コアあたり 128）です。これにより、メモリ転送およびパディングに関する非効率性が排除されるためです。

緩和する方法

メモリに収まり 64 の倍数である最大のバッチサイズを使用することをおすすめします。これを実現する最も簡単な方法は 1024 から開始することであり、メモリ不足エラーが発生した場合は、モデルが正常に実行されるまでバッチサイズを減らしてみます。モデルのバッチサイズを変更するには、同じモデル精度（学習率など）を実現するために他のハイパーパラメータを調整する必要がある場合がありますが、これはケースバイケースで評価する必要があります。

層サイズが小さすぎる

パフォーマンスの問題の説明

モデルで行列乗算または畳み込みが多くを占めていても、入力テンソルが小さい場合、TPU は最も効率的に動作しないことがあります。他のハードウェアと比較すると、TPU はバッチサイズと層サイズの両方が大きい場合に最も効率的に動作します（たとえば、ディメンション >= 512）。

モデルに影響しているかどうかを確認する方法

一般に、層サイズが 128 よりも小さいと、128 が TPU 行列乗算ユニットの固有のディメンションであるため、TPU の効率は悪くなります。全結合層の場合、高い効率を実現するには 512 の最小隠れサイズをおすすめします。畳み込み層は、通常、等しい効率レベルを実現するために、全結合層ほどの大きさである必要はないことに注意してください。たとえば、サイズ 256 の 3 × 3 畳み込みでは、サイズ 2048 の全結合層と比較して同様の（高い）効率が実現されます。3 × 3 × 256 = 2304 であるためです。

緩和する方法

モデル内の層サイズを小さくする主な動機がトレーニング速度である場合は、より大きい層を使用して TPU でモデルを再ベンチマークすることをおすすめします。たとえば、層の出力サイズを 256 から 512 に増やすと、モデルが 2 倍の計算を実行してもトレーニング時間は 20% しか増加しません。

演算レベルのモデルプロファイリング

パフォーマンスのボトルネックを特定するために、演算レベルの実行時間とメモリ使用量を測定することが役に立つことがあります。これを行う方法については、Cloud TPU ツール: Trace Viewer ガイドをご覧ください。

モデルの精度の低下のデバッグ

Cloud TPU エコシステムの目標の 1 つは、現在 CPU または GPU でトレーニングしているモデルを TPU でトレーニングしたときに、バッチサイズや学習率などのハイパーパラメータを多少調整するだけで、同様の精度が実現されることです。ただし、TPU でモデルをトレーニングすると、ユーザーは精度の低下に気付くことがあります。このような問題のデバッグは、ニューラルネットワークトレーニングのランダムな性質のために、非常にもどかしい場合があります。このセクションでは、モデルを TPU に移植する際に、モデルの精度の低下の根本原因を特定する方法について説明します。

データシャーディング（データ並列処理）について

TensorFlow の主要な目標の 1 つは、CPU、GPU、TPU のいずれで実行されるかにかかわらず、各演算によってほぼ同じ結果が生成されることです。ランダム演算など、これにはいくつかの例外があります。一般に、TPU と CPU で非ランダム演算の出力に大きな違いがある場合は、TPU サポートにバグとして報告してください。

ただし、トレーニングパイプライン全体としては、CPU / GPU と TPU でのトレーニングには大きな違いがあります。TPUEstimator と use_tpu=False を使用すると、TensorFlow はその標準の実行エンジンにフォールバックします。このエンジンでは、ステップごとに 1 つのバッチでトレーニングが行われます。ただし、実際の TPU でトレーニングを行う場合は、TensorFlow によってデータシャーディングが実行されます。これは、「同期 SGD によるデータ並列処理」とも呼ばれます。その理由は、各 Cloud TPU は独立したプロセッシングユニットとして動作する 8 個の TPU コアで構成されていることです。そのため、トレーニングのステップごとに、各 TPU コアにデータのバッチが渡され、重みの勾配が計算され、勾配が相互に交換されて、重み更新が計算されます。デフォルトでは、損失はコア全体で平均化されますが、代わりに CrossShardOptimizer のパラメータを変更することで合計できます。

モデルの全損失を、独立したサンプルあたりの損失の平均（または合計）として計算できる場合、この手順は数学的には単一の大きなバッチでのトレーニングと同等です。独立したサンプルあたりではない最も一般的な演算はバッチ正規化であり、これはコアあたりのバッチごとに個別に実行されます。たとえば、バッチサイズの合計が 128 の場合、コアあたりのバッチサイズは 16 であり、8 個のコアのそれぞれが固有の 16 サンプルにわたってバッチ正規化を行います。場合によっては、小さいバッチ（32 未満など）にわたってバッチ正規化を行うことにより、精度が低下することがわかっています。理想的なシナリオでは、TPU でトレーニングするときのバッチサイズの合計を大きくすることができるので（たとえば、256～1024）、そのサイズのバッチは大きな問題になりません。しかし、このようなバッチサイズが大きすぎてメモリに収まらない場合は、シャーディングの影響をケースバイケースで評価する必要があります。

シャーディングによって複雑になっているため、モデルの精度の低下をデバッグするための最初のステップは、決定論的なシングルコア TPU トレーニングを行って、それを CPU / GPU でトレーニングされたモデルと比較することです。一般的に、これは収束までモデルをトレーニングする必要がないため、迅速に行うことができます。

決定論的なトレーニング

モデルの精度の違いをデバッグすることが難しい理由の 1 つは、TensorFlow ではモデルがトレーニングされるたびに、異なる重みの初期化とデータシャッフルが使用されることです。決定論的になるようにトレーニング手順を変更して、複数の実行でほぼ同一のモデルが生成されるようにすると有益です。このセクションでは、MNIST チュートリアルを決定論的に実行する方法を示します。

CPU で 1 ステップ実行することにより、初期チェックポイントファイルを生成します。このステップは、決定論的な重みの初期化を実現するために使用されます。これは変数イニシャライザをシードすることによっても実現できますが、そちらの方が困難です。

# Run training for 1 step to create an initial checkpoint.
python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/init_output \
  --random_seed=12345 \
  --iterations=1
  --train_steps=1

入力関数内のデータシャッフル関数が、ランダムシードを使用するように変更します。これは MNIST チュートリアルですでに行われています。これは、入力データ処理演算に有効です（入力データ処理演算は常に CPU で実行されるため）。モデル関数でのランダム演算は、TPU と CPU の間で決定論的ではない場合があります。次に例を示します。

# In the flag definitions
tf.flags.DEFINE_integer("batch_size", None, "Random seed for training")

# In the input_fn
if FLAGS.random_seed is not None:
dataset = dataset.shuffle(seed=FLAGS.random_seed)

同じモデルを CPU で 2 回実行し、トレーニングが決定論的であることを確認します。トレーニングは妥当なステップ数（たとえば、1000）で実行する必要がありますが、CPU では非常に遅くなる可能性があるため、収束まで実行する必要はありません。

CPU トレーニングはシングルコア TPU トレーニングと比較されるため、単一の TPU コアに収まるバッチサイズを使用します（通常、バッチ全体を 8 で割ったサイズ）。TensorFlow では実行間でビット対ビットの決定論は保証されませんが、損失は非常に近似するはずです。

# Copy the initial weights
gsutil mkdir ${STORAGE_BUCKET}/cpu_output_1
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_1
gsutil mkdir ${STORAGE_BUCKET}/cpu_output_2
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_2

# Run 1
python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/cpu_output_1 \
  --batch_size=128 \
  --random_seed=12345 \
  --train_steps=2000 \
  --eval_steps=10

# Output 1
accuracy = 0.9910644, global_step = 1000, loss = 0.025323588

# Run 2
python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/cpu_output_1 \
  --batch_size=128 \
  --random_seed=12345 \
  --train_steps=2000 \
  --eval_steps=10

# Output 2
accuracy = 0.9910644, global_step = 1000, loss = 0.025323414

シングルコア TPU トレーニング

MNIST チュートリアルを決定論的に実行できたら、次のステップでは、CPU でのトレーニング結果を TPU で複製し、単一の TPU コアを使用して、問題がデータシャーディングに関連するか TPU 実行エンジン自体に関連するかを特定します。

MNIST チュートリアルでシングルコアトレーニングと評価を実行する方法は次のとおりです。

# Use the same weight initialization as the CPU
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/tpu_output

# Run training for 1000 steps
python mnist.py \
    --use_tpu=True \
    --master=$GRPC_SERVER \
    --train_file=${STORAGE_BUCKET}/data/train.tfrecords \
    --model_dir=${STORAGE_BUCKET}/tpu_output \
    --random_seed=12345 \
    --batch_size=128 \
    --num_shards=1 \
    --train_steps=1000 \
    --eval_steps=10

  accuracy = 0.9910644, global_step = 1000, loss = 0.02514153

損失は CPU トレーニング済みモデルと完全には一致しませんが、近似するはずです。ご自身のモデルで近似しない場合は、TPU 実行エンジンでバグが見つかったことを示している可能性があります。バグレポートを TPU サポートに提出する前に、次のことを再確認してください。

num_shards=1 を TPUConfig に渡しています。
モデル関数にランダム演算はなく、入力関数内のランダム演算は正しくシードされています。
CPU と TPU のトレーニングに同じ初期チェックポイントファイルを使用しています。

マルチコア TPU トレーニングのデバッグ

モデルにおいて、CPU とシングルコア TPU で同じ損失が実現した場合、問題は次のいずれかである可能性があります。

（a）低下の原因は、ニューラルモデルを異なる初期化でトレーニングする際の自然でランダムな差異です。

（b）低下の原因は、TPU でのデータシャーディングに関する問題です。

（a）が問題であるかどうかを判断するには、上記のように、重みの初期化を使用して CPU / GPU とマルチコア TPU でモデル全体を再トレーニングすることが有効な場合があります。

精度の低下が統計的に有意であると判断した場合、データシャーディングに関して可能性が最も高い問題は次のとおりです。

モデルで損失をサンプルあたりのエラーの合計として計算する場合は、reduction=losses.Reduction.SUM を CrossShardOptimizer に渡します。デフォルトでは、CrossShardOptimizer は合計ではなく、損失の平均を計算します。
モデルでバッチ正規化を使用する場合、バッチサイズの合計が 256 未満（たとえば、コアあたり 32 未満）であると、精度が低下する可能性があります。
モデルにバッチ式の損失関数がある場合、これはシャーディングの影響を受けます。そのような損失関数は、通常、非常に特殊化されています。たとえば、Karras et al.2017 は、敵対的生成ネットワークをトレーニングする際に、バッチ弁別子を使用します。

利用可能な機能に関するよくある質問

TPU を推論に使用できますか？

はい。TPU はトレーニングと推論の両方に使用できます。たとえば、ResNet チュートリアルでは、トレーニングループ中に定期的な評価を行います。モデルサービングの場合、注意すべき点がいくつかあります。特に、TPU ソフトウェアスタックは現在、レイテンシではなくスループットに最適化されています。入力の 1 つのバッチに対して推論を実行し、結果を待機すると、現在は少なくとも 10 ミリ秒のオーバーヘッドがあります。これは、低レイテンシでのサービングでは問題となる可能性があります。

このオーバーヘッドは、今後の TensorFlow リリースで大幅に削減される予定です。

TPU で利用できない組み込み TensorFlow 演算はありますか？

少数の組み込み TensorFlow 演算は、現在 TPU で利用できません。現在の回避策の詳細については、利用可能な TensorFlow 演算のガイドをご覧ください。

TPU のカスタム演算を記述するにはどうすればよいですか？

TPU で実行される TensorFlow 演算は、XLA HLO で実装されています。これは、少数の低レベルの関数を使用して高レベルのテンソル演算を定義するための言語です。XLA は TensorFlow のオープンソースリリースに含まれているので、HLO で演算を記述することが技術的に可能です。既存の実装の大部分は、tf2xla ディレクトリにあります。ただし、これは任意の C++ または Python コードではなく、TPU での限られたテンソル演算の実行のみが可能です。HLO で実装できる最も一般的なテンソル演算は、すでに記述されています。TensorFlow の今後のリリースでは、TPU のトレーニング / 推論中に標準的な CPU 演算を効率的に実行する機能がサポートされる予定です。

プレースホルダとフィード辞書を TPU で使用できますか？

この使用パターンは TPU で技術的に可能ですが、単一の TPU コアのみが使用され、過剰なオーバーヘッドが発生するため、使用しないことを強くおすすめします。代わりに、トレーニングパイプラインを作成するには、TPUEstimator API と Dataset API を使用します。TPUEstimator と Dataset を使用して単純なトレーニングループを作成する方法の例については、ResNet チュートリアルをご覧ください。

強化学習（RL）モデルを TPU でトレーニングできますか？

強化学習は幅広い技術をカバーしていますが、そのうちのいくつかは、現在は TPU のソフトウェア抽象化と互換性がありません。一部の強化学習構成では、トレーニングループの一部として CPU を使用してブラックボックス「シミュレーション環境」を実行する必要があります。経験的に、これらは TPU に追いつくことができず、重大な非効率をもたらすことになります。TensorFlow の今後のリリースには、「オフポリシー」強化学習を簡単にするための抽象化が含まれる予定です。

TPU で単語埋め込みを使用できますか？

はい。TPU では tf.nn.embedding_lookup() がサポートされています。これは単に、TPU で実装されている tf.gather() のラッパーです。しかし、TPU では tf.nn.embedding_lookup_sparse() はサポートされていません。トレーニング中、tf.embedding_lookup() への入力 ID テンソルは静的形状である必要があることに注意してください（つまり、バッチサイズとシーケンス長はすべてのバッチで同じでなければなりません）。これは、TPU を使用する場合に、すべてのテンソルに対する、より一般的な制限です。

TPU で可変長シーケンスを使用できますか？

TensorFlow で可変長シーケンスを表すには、パディング、tf.while_loop()、推定されるテンソル次元、バケット化など、いくつかの方法があります。現在の TPU 実行エンジンでは、これらのサブセットのみがサポートされています。可変長シーケンスは、tf.while_loop()、tf.dynamic_rnn()、バケット化、パディング、シーケンスの連結を使用して実装する必要があります。

TPU で反復性ニューラルネットワーク（RNN）をトレーニングできますか？

特定の構成では、tf.static_rnn() および tf.dynamic_rnn() は現在の TPU 実行エンジンと互換性があります。より一般的に言うと、TPU では tf.while_loop() と TensorArray の両方がサポートされており、これらは tf.dynamic_rnn() を実装するために使用されます。CuDNN などの特殊なツールキットは、GPU 固有のコードが含まれているため、TPU ではサポートされていません。TPU で tf.while_loop() を使用するには、TPU 実行エンジンがメモリ使用量を静的に決定できるように、ループ反復回数の上限を指定する必要があります。

TPU で敵対的生成ネットワーク（GAN）をトレーニングできますか？

GAN のトレーニングは、通常、ジェネレータのトレーニングと弁別子のトレーニングを頻繁に交互に切り替える必要があります。現在の TPU 実行エンジンでは、1 つの実行グラフのみがサポートされます。グラフを交互に切り替えるには完全な再コンパイルが必要で、これには 30 秒以上かかる場合があります。この制限は、今後の TensorFlow リリースで改善される予定です。

可能性のある 1 つの回避策は、ジェネレータと弁別子の両方の損失の合計を常に計算するが、これらの損失に 2 つの入力テンソル g_w と d_w を乗算することです。ジェネレータをトレーニングするバッチでは、g_w=1.0 と d_w=0.0 を渡すことができ、弁別子をトレーニングするバッチについて、逆の場合も同じです。

TPU でマルチタスク学習モデルをトレーニングできますか？

タスクを集約損失関数を持つ 1 つの大きなグラフとして表すことができる場合、マルチタスク学習に特別なサポートは必要ありません。ただし、TPU 実行エンジンでは現在、1 つの実行グラフのみがサポートされています。したがって、変数を共有するが構造が異なる複数の実行グラフをすばやく交互に切り替えることはできません。実行グラフを変更するには、グラフコンパイルステップの再実行が必要で、これには 30 秒以上かかる場合があります。

TPU では eager モードはサポートされますか？

いいえ。eager モードでは新しい動的実行エンジンが使用されますが、TPU では XLA が使用されます。XLA では実行グラフの静的コンパイルが実行されます。

TPU ではモデルの並列処理はサポートされますか？

モデルの並列処理（つまり、単一の TPU デバイスの複数のコアで、同一でない TPU プログラムを実行すること）は、TPU で現在はサポートされていませんが、今後の TensorFlow リリースでサポートされる予定です。

TPU の中間テンソルの実際の値を調べるにはどうすればよいですか（tf.Print や tfdbg と同様に）？

この機能は現在、TPU ではサポートされていません。TPU での開発向けの推奨パターンは、TPUEstimator フレームワークを使用してモデルを実装することです。このフレームワークでは、use_tpu フラグを使用して TPU と CPU / GPU 間で簡単に移行できます。標準の TensorFlow ツールを使用して CPU / GPU でモデルをデバッグし、モデルで本格的なトレーニングの準備ができたら TPU に切り替えることをおすすめします。

トレーニングスキームが、TPUEstimator API には複雑または専門的すぎます。使用できる低レイヤ（低水準）の API はありますか？

TPUEstimator は、Cloud TPU での TPU トレーニングの主要なフレームワークです。ただし、TPUEstimator によってオープンソースの TensorFlow の一部である tpu API がラップされているので、低レベルの tpu API を直接使用することは技術的には可能です（サポートはされていません）。トレーニングパイプラインで TPU と CPU の間の頻繁な通信が必要な場合や、実行グラフを頻繁に変更する必要がある場合、TPU で計算を効率的に実行できません。TensorFlow の今後のリリースでは、両方の機能が向上する予定です。

概要

TPU サーバーへの接続の問題

一般的なエラーのデバッグ

ローカル ファイルシステムを使用できない

サポートされないデータ型

動的形状がサポートされない

利用できない TensorFlow 演算

メモリ不足のエラー メッセージ

CrossShardOptimizer を使用していない

TPU サーバーに接続できない

トレーニング中のエラー

実行の停止を起こす問題

メモリ使用量の削減

モデルの重みの数が多い

テンソル パディングが過剰である

バッチサイズが大きすぎる

モデルが大きすぎる

トレーニング速度の向上

使用していない TPU コアがある

ループあたりの反復回数が少なすぎる

入力処理のボトルネック

非行列乗算演算が多すぎる

テンソル パディングが過剰である

バッチサイズが小さすぎる

層サイズが小さすぎる

演算レベルのモデル プロファイリング

モデルの精度の低下のデバッグ

データ シャーディング（データ並列処理）について

決定論的なトレーニング

シングルコア TPU トレーニング

マルチコア TPU トレーニングのデバッグ

利用可能な機能に関するよくある質問

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

ローカルファイルシステムを使用できない

メモリ不足のエラーメッセージ

`CrossShardOptimizer` を使用していない

テンソルパディングが過剰である

テンソルパディングが過剰である

演算レベルのモデルプロファイリング

データシャーディング（データ並列処理）について