TensorFlow ワークロードをプロファイリングする

Cloud TPU で TensorFlow モデルをプロファイリングするには、TensorBoard と TPU TensorBoard プラグインを使用します。TensorBoard は TPU VM にプリインストールされています。TPU TensorBoard プラグインをインストールしてパフォーマンス プロファイルをキャプチャする方法については、Cloud TPU ツールでモデルをプロファイリングするをご覧ください。Cloud TPU の一般的なパフォーマンス情報については、Cloud TPU パフォーマンス ガイドをご覧ください。

詳しくは、TensorBoard のコールバックをご覧ください。

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

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

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

        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

  • XLA コンパイラは、バッチ正規化の TensorFlow の融合バリアントを効率的に減らすことができます。これらを使用する方が、別の方法を使用するよりかなり効率的です。

    • tf.nn.batch_normalization よりも tf.nn.fused_batch_norm を優先します。
    • tf.layers.batch_normalization の場合は、「fused」引数を true に設定します。