トラブルシューティング

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

概要

Cloud TPU で TensorFlow モデルを実行する際のおすすめの方法は、TPUEstimator API を使用することです。TensorFlow の Estimator API を現在使用している場合、通常は数行のコードを変更するだけで TPUEstimator に切り替えられます。TPUEstimator にデータを読み込む際のおすすめの方法は、tf.data API を使用することです。TPUEstimator と tf.data を実装する方法の実例については、Cloud TPU での TPUEstimator API の使用にある MNIST の例をご覧ください。MNIST Estimator から TPUEstimator へもご覧ください。

モデルを TPUEstimator に変換した後、そのモデルがフラグ use_tpu=False で動作することを確認します。このフラグを false に設定すると、TensorFlow は Estimator API にフォールバックし、TPU に関連するコードは使用されません。use_tpu=False でモデルを実行しているときに発生する問題は TPU に関係なく、このガイドの範囲外です。代わりに、TensorFlow プログラマーガイドをご覧ください。

理論的には、TPUEstimator と use_tpu=False を使用してモデルを正常に実行できるようになった後は、use_tpu=True を設定して、master で TPU サーバー URL を指すだけで（通常はクラスタリゾルバを使用）、そのモデルを TPU 上で実行できます。ただし、TensorFlow モデルは非常に複雑な場合があり、TPU では独自の実行エンジンが使用されるため、TPU 固有の問題が発生する場合があります。このような問題は、次の大きなカテゴリに分類されます。

トレーニングスクリプトが 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  READY
```
--tpu に正しい値を渡していること（上記の例では 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 stop $TPU_SERVER_NAME && gcloud compute tpus start $TPU_SERVER_NAME を使用して手動でサーバーを再起動します。上記の例では、$TPU_NAME は demo-tpu です。この処理には数分かかることがあります。
上記の ... tpus list コマンドを再実行し、TPU が READY 状態になるのを待ちます。この処理には数分かかることがあります。
MNIST チュートリアルを再度実行してみます。
この手順に従っても MNIST チュートリアルを正常に実行できない場合は、サポートの利用で説明しているメカニズムのいずれかを使用してサポートをご依頼ください。

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 サーバー上で実行されることにご注意ください。Cloud Storage で TPU との使用を適切に構成する方法については、Cloud Storage バケットへの接続のガイドをご覧ください。

tf.data.Dataset.cache() でローカルファイルシステムへのキャッシュ保存ができない

エラーメッセージ

tensorflow.python.framework.errors_impl.UnimplementedError: File system scheme '[local]' not implemented (file: '[filename].lockfile')

詳細

tf.data.Dataset はキャッシュに保存できます。.cache() 呼び出しには、次の 2 つの実装があります。

メモリ内 - 引数が渡されない場合。
ファイルシステム上 - 引数としてファイルパスが渡された場合。

Cloud TPU では、(1) は（利用可能なメモリに収まる限り）機能しますが、ローカルファイルシステムに保存するときには (2) は機能しないため、上記のエラーが発生します。

以下のコードスニペットは、この 2 つの状況を示しています。

(1)
 import tensorflow as tf

def main():
  print('Hello world!')
  ds = tf.data.Dataset.range(10)
  ds = ds.cache()

runs to completion.

(2)
 import tensorflow as tf

def main():
  print('Hello world!')
  ds = tf.data.Dataset.range(10)
  ds = ds.cache('/tmp/foo')

generates the error.

API ガイドには、tf.data.Dataset.cache() に関する詳しい説明があります。

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

エラーメッセージ

TypeError: DataType is not a supported TPU infeed type.

詳細

現在、tf.float32、tf.int32、tf.bfloat16、tf.bool データ型のみが TPU でサポートされています。その他の一般的なデータ型（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 サーバーに再接続しようとして DeadlineExceededError のような新しいエラーが発生した場合は、コマンド gcloud compute tpus stop $TPU_SERVER_NAME && gcloud compute tpus start $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 の倍数である。

または
バッチサイズの合計は 1,024 の倍数（TPU コアあたり 128）であり、特徴ディメンションは 8 の倍数である。

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

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

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

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

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

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

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

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

モデルが大きすぎる

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

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

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

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

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

トレーニング速度の向上

モデルを TPU で正常に実行できるが、トレーニング速度が想定ほど速くない場合があります。このセクションでは、速度の向上の実現に役立ついくつかの方法の概略を説明します。

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

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

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 ツール: 入力パイプライン分析ツールの指示に従います。

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

緩和する方法

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

緩和する方法

モデル内の行列乗算の不足が、主に他のハードウェアでのトレーニング速度の問題が起因となった場合は、速度のパフォーマンス向上のために、それらのモデルを 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 ツール: トレースビューアのガイドをご覧ください。

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

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

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

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

ただし、トレーニングパイプライン全体としては、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 回実行し、トレーニングが決定論的であることを確認します。トレーニングは妥当なステップ数（たとえば、1,000）で実行する必要がありますが、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 \
    --train_steps=1000 \
    --eval_steps=10

  accuracy = 0.9910644, global_step = 1000, loss = 0.02514153

損失は CPU トレーニング済みモデルと完全には一致しませんが、近似するはずです。対象のモデルで近似しない場合は、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 サーバーへの接続の問題

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

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

tf.data.Dataset.cache() でローカル ファイル システムへのキャッシュ保存ができない

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

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

利用できない TensorFlow 演算

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

CrossShardOptimizer を使用していない

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

トレーニング中のエラー

実行の停止を起こす問題

メモリ使用量の削減

モデルの重みの数が多い

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

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

モデルが大きすぎる

トレーニング速度の向上

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

入力処理のボトルネック

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

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

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

層サイズが小さすぎる

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

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

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

決定論的なトレーニング

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

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

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

tf.data.Dataset.cache() でローカルファイルシステムへのキャッシュ保存ができない

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

`CrossShardOptimizer` を使用していない

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

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

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

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