パフォーマンス ガイド

Cloud TPU は低コストで高パフォーマンスを提供します。Cloud TPU のパフォーマンスをさらに向上させるには、アプリケーションの Cloud TPU 構成パラメータを調整し、パフォーマンスを制限するボトルネックを特定して解決します。

このガイドでは、Cloud TPU のパフォーマンスを最大限に高めるために、次の方法について説明します。

また、次のリソースでは、以下の方法について説明します。

モデル処理のパフォーマンス

このセクションでは、モデルのパフォーマンスを低下させる可能性がある一般的な問題と、その処理方法について説明します。

  1. データの前処理にかかる時間が長すぎる

    TensorFlow-TPU ソフトウェア スタックにより、ユーザーは TPU にデータをフィードする前に、CPU 上で複雑なデータの前処理を実行できます。TPU は非常に高速であるため、入力データが非常に大規模である場合、または複雑な方法で処理されている場合、入力データ処理がボトルネックになる可能性があります。

    Cloud TPU プロファイリング ツールにより、入力処理がモデルのボトルネックかどうかを簡単に測定できます。この場合、高額なオフラインの前処理は、すべてのトレーニングのすべてのエポックでそのコストが発生しないように、1 回限りのコストとして実行します。

  2. シャーディング（コア間でのバッチの分割）のためにバッチサイズが小さすぎる

    n のトレーニング バッチサイズを使用している場合、これは 8 個の TPU コアに分割またはシャーディングされます。 n = 128 が TPU を単一の P100 と比較する場合、GPU はピーク FLOP の 70% で実行される可能性がありますが、TPU はピーク FLOP の 20% でしか実行されません（シャーディングされた TPU バッチサイズは 128 ではなく実際には 16（128 / 8）であるため）。この例では、TPU のバッチサイズが 1,024 の場合、TPU メモリをより活用できます。

    メモリ使用量を最適化するには、メモリに収まる最大サイズのバッチサイズを使用します。各 TPU コアは、128 x 128 メモリセル マトリックスを処理します。一般に、TPU メモリを最大限効果的に使用するには、128 で割り切れるバッチサイズである必要があります。

XLA コンパイラのパフォーマンス

XLA は、TPU、CPU、GPU などのプラットフォーム用のバイナリを生成できる機械学習用のコンパイラです。これは標準の TensorFlow コードベースの一部です。 Cloud TPU の TensorFlow モデルは XLA グラフに変換され、XLA は TPU 実行可能ファイルにコンパイルします。 XLA と TensorFlow の相互作用の詳細については、XLA の概要をご覧ください。

Cloud TPU ハードウェアは、CPU と GPU とは異なります。CPU は、高パフォーマンスのスレッドの数が少ないという特徴があります。GPU は、低パフォーマンスのスレッドの数が非常に多いという特徴があります。Cloud TPU には 128 x 128 のマトリックス ユニットがあるため、1 サイクルにつき 16 K 演算を実行できる 1 つの非常に強力なスレッド、またはパイプライン方式で接続された 128 x 128 の小さい単純なスレッドと考えることができます。これに対応して、メモリのアドレスを指定する場合は、8 の倍数（浮動小数点数）が適しており、マトリックス ユニットを対象とする演算には 128 の倍数が適しています。

タイリングの結果

Cloud TPU 内の配列はタイル状になっています。これは、ディメンションの 1 つを 8 の倍数にパディングし、別のディメンションを 128 の倍数にパディングする必要があります。 XLA はデータ レイアウトの変換を行い、ハードウェアがデータを効率的に処理できるようにデータをメモリに配置します。 これらの変換はヒューリスティクスによって行われます。これらはほとんどの場合有益ですが、コンパイラが間違った処理を行う可能性は常にあります。最高のパフォーマンスを得るには、さまざまなモデル構成を試すことが有益です。

経験則: パディングのコストを最小限に抑える

タイル状のメモリスキームを使用した結果は、効率的なメモリ使用率はパディングのオーバーヘッドで無駄に消費されたメモリ量に依存するというものです。 最も効率的な方法で Cloud TPU を使用するには、タイリングのオーバーヘッドを最小限に抑えるディメンション サイズを使用します。

たとえば、畳み込み結果には、（1）バッチ、（2）出力フィーチャ、（3）出力空間のディメンションがあります。バッチ ディメンションまたは出力フィーチャ ディメンションの一方は 8 の倍数にパディングされますが、もう一方は 128 の倍数にパディングされます。出力空間ディメンションはパディングされません。

一般的に、空間ディメンション（tf.nn.pool、tf.conv2d など）を使用した演算では、空間ディメンションはパディングされません。

経験則: バッチ ディメンションとフィーチャ ディメンションの効果的な値を選択する

バッチ ディメンションとフィーチャ ディメンションはパディングの対象となるため、バッチサイズとフィーチャ サイズを決定する際は注意が必要です。128 x 128 のマトリックス ユニットを最大限に活用するには、バッチまたはフィーチャに適した大きな値（>=128）、できれば両方の値を求めます。

ほとんどの場合、バッチサイズ 128 は Cloud TPU マトリックス ユニットを完全に占有するのに十分です。より大きなバッチサイズで実行しても Cloud TPU でうまく機能しますが、128 の倍数であるバッチサイズを使用することをおすすめします。このような構成でモデルが機能しない場合は、パディングの影響を最小限に抑えるために、8 の倍数であるバッチサイズを使用してみてください。

同様に、フィーチャ ディメンションも、XLA レイアウト アルゴリズムによる決定に応じて、8 または 128 のいずれかにパディングされるディメンションにマッピングされます。つまり、フィーチャ ディメンションは、8 または 128 の倍数で最小のスペースを浪費します。

融合

融合は、XLA コンパイラがプログラムの最適化に使用する一般的な手法です。 融合演算は、組み合わせて実行する、複数の構成演算の組み合わせです。

たとえば、次の一連の演算について考えてみます。

tmp = tf.add(x, y)
result = tf.multiply(tmp, z)

このコードは、次の擬似コードとほぼ同じです。

for (i = 0; i < element_count; i++) {
  tmp[i] = x[i] + y[i];
}

for (i = 0; i < element_count; i++) {
  result = tmp[i] * z[i];
}

融合により、配列アクセスが同時に発生します。

for (i = 0; i < element_count; i++) {
  result = (x[i] + y[i]) * z[i];
}

この例では、メモリのラウンド トリップ数が削減され、XLA は「tmp」にスペースを割り当てる必要がありませんでした。

融合は非常に重要な最適化であり、次のいくつかの点で Cloud TPU にメリットがあります。

  • 低速のメインメモリに中間結果を保存する必要をなくして、メモリ転送を低減します。
  • 通常は使用されないようなハードウェア ユニットを有効に活用できます。
  • 同時にアクティブ化されるバッファの数が少ないため、モデルのメモリ使用率を低減できます。

ブロードキャスト

ブロードキャストは、異なるが互換性のあるシェイプを持つ 2 つのテンソルが組み合わされたときに暗黙的に発生します。

たとえば、tf.add(vector, matrix) は、マトリックスのシェイプにベクトルをブロードキャストすることを要求します。演算の結果は、マトリックスと同じシェイプになります。詳細については、配列のブロードキャストのガイドをご覧ください。

ブロードキャストは多くの場合コンシューマと融合できますが、ブロードキャストを実体化するとパフォーマンスが低下し、メモリの負荷が増加します。

次の例では、ベクトルとマトリックスを追加する際の暗黙のブロードキャストを、ブロードキャストを実体化する argmax と融合することはできません。

tf.argmax(tf.add(vector, zero_matrix), axis=0)

TensorFlow の関数固有の推奨事項

Cloud TPU で使用可能な TensorFlow 演算の完全なリストをご覧ください。

tf.matmul

  • いずれかのオペランドの結果の転置は事実上自由です。
  • tf.matmul はその入力と出力への融合をサポートしていることに注意してください。このため、tf.matmul の出力に直接適用される活性化関数やバイアスのオーバーヘッドが少なくなります。

tf.nn.conv_n_d、tf.nn.depthwise_conv2d、tf.nn.separable_conv2d

  • 活性化では、バッチ ディメンションと特徴ディメンションが 8 または 128 の倍数にパディングされます。
    • まず XLA は、モジュール内で畳み込みに最も一般的なバッチ ディメンションのサイズを追跡します。これにより、順方向畳み込み、活性化勾配畳み込み、カーネル勾配畳み込みを区別できます。
    • 最も一般的なバッチサイズが 64 以上の場合:
      • 順方向畳み込みと逆方向畳み込みの場合、バッチは 128 の倍数にパディングされ、特徴は 8 の倍数にパディングされます。
      • 勾配更新畳み込みの場合、バッチは 8 の倍数にパディングされ、特徴は 128 の倍数にパディングされます。
    • 最も一般的なバッチサイズが 64 未満の場合:
      • 順方向畳み込みと逆方向畳み込みの場合、バッチは 8 の倍数にパディングされ、特徴は 128 の倍数にパディングされます。
      • 勾配更新畳み込みの場合、バッチは 128 の倍数にパディングされ、特徴は 8 の倍数にパディングされます。
      • 転置で入力する特徴ディメンションとバッチ ディメンションを入れ替えるだけの場合は、畳み込みに送信する直前に活性化を自由に転置できます。
  • カーネルでは、入出力の特徴ディメンションは 8 または 128 の倍数にパディングされます。その決定は、カーネルのプロデューサーや他のコンシューマの影響を受けます。
    • 転置で入力と出力の特徴ディメンションを入れ替えるだけの場合は、畳み込みに送信する直前のカーネルを自由に転置できます。
  • 結果では、バッチ ディメンションと特徴ディメンションは 8 または 128 の倍数にパディングされます。
    • 転置でバッチ ディメンションと出力特徴ディメンションを入れ替えるだけの場合は、畳み込みの結果を自由に転置できます。
  • tf.nn.conv_n_d は、結果、活性化、カーネルへの融合をサポートしていることに注意してください。このため、出力に直接適用される活性化関数やバイアスのオーバーヘッドが少なくなります。

tf.nn.avg_pool、tf.nn.max_pool

  • パディング ルールが適用されます。空間ディメンションはバッチや特徴よりも重要です。バッチと特徴はそれぞれ、8 または 128 の倍数にパディングされます。
  • 通常、プール演算のレイアウトは、そこに流れて出入りする畳み込みと一致します。
  • tf.nn.max_pool の勾配計算は、同等の tf.nn.avg_pool より低速になる場合があります。可能な場合は、max-pooling から average-pooling への変更を検討してください。

tf.concat、tf.slice、tf.strided_slice

  • 不要なスライスや連結は避けてください。パディングされたディメンション内のスライスと連結は、高コストになります。
    • スライス ディメンションにパディングのオーバーヘッドがなければ、データの移動が最小限に抑えられます。

tf.transpose

  • tf.matmul またはその結果のオペランドのいずれかの転置は自由に行うことができます。
  • バッチ ディメンションと入力の特徴ディメンションを入れ替える場合は、tf.conv_n_d の活性化を自由に転置できます。
  • 入力の特徴ディメンションと出力の特徴ディメンションを入れ替える場合は、tf.conv_n_d のカーネルを自由に転置できます。
  • バッチ ディメンションと出力の特徴ディメンションを入れ替える場合は、tf.conv_n_d の結果を自由に転置できます。

tf.batch_to_space、tf.space_to_batch、tf.space_to_depth、tf.depth_to_space

  • これらは、パディングされたディメンションからパディングされていないディメンション（またはその逆）へのデータの移動が必要なため、コストがかかります。

tf.reshape

  • パディングされたディメンション内のデータを移動すると、Cloud TPU でのリシェイプが高コストになる場合があります。
  • かなりのパディングが存在する場合、ホスト上の R1 にデータをリシェイプし、デバイス上高位のディメンションのシェイプにリシェイプすると有益です。これにより、ホストと端末の間の転送をより効率的に行うことができます。
    • また、パッケージ化されたパラメータをオンデマンドで展開できるため、ピーク時のメモリ使用率の向上にも役立ちます。

tf.random_uniform、tf.distributions.Bernoulli、tf.random_normal、tf.multinomial

  • 一様分布またはベルヌーイ分布の擬似乱数生成は非常に高速です。
  • 正規分布は、一様分布またはベルヌーイ分布よりも少しコストがかかります。
  • カテゴリ分布と多項分布の擬似乱数生成はかなりコストがかかります。

tf.reduce_all、tf.reduce_any、tf.reduce_logsumexp、tf.reduce_max、tf.reduce_min、tf.reduce_prod、tf.reduce_sum

  • 入力シェイプと出力シェイプが同じ場合、融合によって複数の縮約を並行して実行できます。
    • 可能な場合は、縮約の順次連鎖を平行連鎖に書き換えてみてください。

  • 縮約は要素ごとの演算を入力に融合させますが、出力には融合させません。可能な場合は、式を書き換えて融合を進めます。 例:

        tf.multiply(tf.reduce_sum(x), y)

    以下のようになります。

        tf.reduce_sum(tf.multiply(x, y))

tf.nn.batch_normalization、tf.nn.fused_batch_norm、tf.layers.batch_normalization

  • Cloud TPU コンパイラは、バッチ正規化の TensorFlow の融合バリアントを効率的に減らすことができます。これらを使用する方が、別の方法を使用するよりかなり効率的です。
    • tf.nn.batch_normalization よりも tf.nn.fused_batch_norm を優先します。
    • tf.layers.batch_normalization の場合は、「fused」引数を true に設定します。

tf.nn.depthwise_conv2d、tf.nn.separable_conv2d

  • これらはまだ完全には最適化されておらず、同等の通常の畳み込みよりもパフォーマンスが劣る可能性があります。