パフォーマンス ガイド

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

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

• XLA コンパイラの効率を向上させます。
• 調整可能な TensorFlow の関数セットを特定して使用します。

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

• 入力パイプラインを調整します。
• Cloud TPU ツールを使用して、パフォーマンスのボトルネックを特定します。
• より速く収束するように全体的なパフォーマンスを向上させます。

注: これらのページでは一連のガイドラインのみを示します。いくつかの理由（基本的なモデル アーキテクチャの要件など）でアドバイスから逸脱する必要がある場合があります。これらのガイドラインは、Cloud TPU ソフトウェア スタックの改善に伴って変更されます。

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

• 活性化では、バッチ ディメンションとフィーチャ ディメンションは 8 または 128 の倍数にパディングされます。
  • まず XLA は、モジュール内の畳み込みのために最も一般的なバッチ ディメンションのサイズを追跡します。これにより、順方向畳み込み、活性化勾配畳み込み、カーネル勾配畳み込みを区別できます。
  • 最も一般的なバッチサイズが 64 以上の場合:
    • 順方向畳み込みと逆方向畳み込みの場合、バッチは 128 の倍数にパディングされ、フィーチャは 8 の倍数にパディングされます。
    • 勾配更新畳み込みの場合、バッチは 8 の倍数にパディングされ、フィーチャは 128 の倍数にパディングされます。
  • 最も一般的なバッチサイズが 64 未満の場合:
    • 順方向畳み込みと逆方向畳み込みの場合、バッチは 8 の倍数にパディングされ、フィーチャは 128 の倍数にパディングされます。
    • 勾配更新畳み込みの場合、バッチは 128 の倍数にパディングされ、フィーチャは 8 の倍数にパディングされます。
    • transpose が入力フィーチャ ディメンションとバッチ ディメンションをスワップするだけの場合は、畳み込みに送信する直前の活性化の転置は自由に行うことができます。
• カーネルでは、入力フィーチャ ディメンションと出力フィーチャ ディメンションは 8 または 128 の倍数にパディングされます。決定は、カーネルのプロデューサーや他のコンシューマの影響を受けます。
  • transpose が入力フィーチャ ディメンションと出力フィーチャ ディメンションをスワップするだけの場合は、畳み込みに送信する直前のカーネルの転置は自由に行うことができます。
• 結果では、バッチ ディメンションとフィーチャ ディメンションは 8 または 128 の倍数にパディングされます。
  • transpose がバッチ ディメンションと出力フィーチャ ディメンションをスワップするだけの場合は、畳み込みの結果の転置は自由に行うことができます。
• 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

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