TPU ポッドでのトレーニング

概要

Cloud TPU のメリットの 1 つは、異なるサイズの TPU タイプ間でトレーニングをスケーリングできることです。TPU タイプはすべて、同じデータ並列アーキテクチャを使用します。唯一異なる点は、並列処理が 8 コアから 2048 コアに増えることです。

多数の TPU を最大限に活用するには、いくつかのトレーニング タスク パラメータを調整する必要があります。このドキュメントでは、共通の問題と、モデルに加える必要がある変更について説明します。

バッチサイズとトレーニング ステップのスケーリング

規模が大きい TPU タイプで線形スケーリングを行うには、コアごとのバッチサイズを同じにします。

たとえば、v2-8 でバッチサイズ 1024 を使用する場合、v2-32 ではバッチサイズ 4096(4 * 1024)を使用します。これは TPU ハードウェアを十分に活用しています。小さいバッチサイズを使用できますが、そうしてもトレーニングを直線的にスケーリングすることはできません。

モデルの多くには、1 つのステップが 1 つのデータバッチに対応する train_steps フラグが含まれています。バッチサイズを増やす場合は、トレーニングのステップ数を減らしてスケーリングして、トレーニング サンプルの合計数が同じになるようにします。

たとえば、元の構成で 1,000 ステップに対するバッチサイズが 1024 であった場合に、バッチサイズを 4,096 に増やしたら、ステップを 250(1000 / 4)に減らすことができます。モデルで epochs フラグを使用する場合、ステップ数をスケーリングする必要はありません。

バッチサイズを大きくすると、モデルの収束動作が変わるため、一部のハイパーパラメータの調整も必要になることがあります。

大規模なバッチ オプティマイザーの使用

スケールアップしてより大きいバッチサイズを使用する場合、規模の大きいバッチサイズに対応するオプティマイザーの使用が必要になることもあります。たとえば、Resnet リファレンス モデルでは、規模の大きい TPU タイプに LARSOptimizer が使用されます。

TPU Pod と同じリージョン内でのリージョン Google Cloud Storage バケットの使用

一般に、TPU トレーニングにおすすめの方法は、常に同じリージョン内のリソースを使用することです。Google Cloud Storage からのデータ転送率がより高くなる傾向があるため、TPU Pod ポッドを使用する場合は、リソースのリージョンが特に重要です。データセットとチェックポイントのトレーニングに使用する TPU と同じリージョンで、リージョンの Google Cloud Storage バケットを使用していることを確認してください。

TPUClusterResolver に対する TPU 名とリージョンの指定

単一デバイスのトレーニングでは、TPU 名または IP アドレスのいずれかを指定できます(例: grpc://1.2.3.4:8470)。

TPU Pod の場合、TPU 名を使用して、TensorFlow がトレーニング分散に利用可能なホストすべての IP アドレスを検出できるようにする必要があります。

TPU Pod の IP アドレスを指定すると、次のエラーが発生します。

ValueError: TPUConfig.num_shards is not set correctly ....

モデルと TPU を別のゾーンやリージョンで実行する必要がある場合は、TPUClusterResolver に TPU ゾーンを指定する必要があります。

tpu フラグと tpu_zone フラグをリファレンス モデルに渡すことも、TPUClusterResolver の値を直接設定することもできます。例:

cluster = tf.contrib.cluster_resolver.TPUClusterResolver(
  tpu='tpuname',
  zone='europe-west4-a',
  project='myproject')

TPUEstimator の使用時に発生する可能性のある関連問題は、クラスタ リゾルバ内で tf.contrib.tpu.RunConfig に渡されません。これは、RunConfig が master パラメータで初期化されている場合に発生することがあります。このため、初期化には cluster パラメータを使用してください。

config = tf.contrib.tpu.RunConfig(cluster=cluster)

TPUEstimator を使用せずに TPU Pod と tf.Session を使用する場合、セッション構成を以下のように構成する必要があります。

with tf.Session(cluster.get_master(), config=config.session_config) as sess:

TPU ポッドでの評価

TPU Pod ポッドでは、評価はサポート対象ではなく、コスト効率も良くありません。評価は、単一デバイスの TPU で mode フラグを使用して実行することをおすすめします。この方法は、大部分のモデルでサポートされています。

TPU Pod または単一デバイス TPU でのトレーニングについては、フラグを train に設定します。評価の場合、単一デバイス TPU に対してのみ、eval を設定します。単一デバイス TPU ではモード train_and_eval も使用できますが、TPU Pod では使用できません。