今日のローカルTensorFlowEverywhereイベントの出欠確認!

Grappler による TensorFlow グラフ最適化

TensorFlow.org で表示 Google Colab で実行 GitHub でソースを表示 ノートブックをダウンロード

概要

TensorFlow は、グラフ実行と Eager Execution の両方を使用して計算を実行します。tf.Graphには、計算の単位を表すtf.Operationオブジェクト(演算)と、演算間を流れるデータの単位を表すtf.Tensorオブジェクトのセットが含まれています。

Grappler は、TensorFlow ランタイム内のデフォルトのグラフ最適化システムです。Grappler は、グラフモード(tf.function内)で最適化を適用して、グラフの簡略化や、プロシージャ間の最適化を可能にする関数本体のインライン化などその他の高レベルの最適化により、TensorFlow 計算のパフォーマンスを向上させます。tf.Graphを最適化すると、グラフノードの計算リソースへのマッピングを最適化することによってデバイスのピークメモリ使用量が削減され、ハードウェアの使用率が向上します。

tf.Graphの最適化をより細かく制御するには、tf.config.optimizer.set_experimental_options()を使用します。

使用可能なグラフオプティマイザ

Grappler はMetaOptimizerと呼ばれるトップレベルのドライバを通してグラフの最適化を実行します。TensorFlowでは、次のグラフオプティマイザの使用が可能です。

  • 定数折り畳みオプティマイザ - 可能な場合にグラフの定数ノードを折り畳むことによりテンソルの値を静的に推測し、定数を使用して結果をマテリアライズします。
  • 算術オプティマイザ - 一般的な副次式を消去し、算術ステートメントを簡略化することにより、算術演算を簡略化します。
  • レイアウトオプティマイザ - テンソルのレイアウトを最適化することにより、畳み込みなどのデータフォーマットに依存する演算をより効率的に実行します。
  • リマッパーオプティマイザ - 一般的に発生するサブグラフを最適化された結合モノリシックカーネルで置き換えることにより、サブグラフをより効率的な実装に再マッピングします。
  • メモリオプティマイザ - グラフを分析して各演算のピークメモリ使用量を確認し、CPU-GPU メモリコピー演算を挿入してGPUメモリをCPUのメモリにスワップすることにより、ピークメモリ使用量を減らします。
  • 依存オプティマイザ - 制御の依存関係を削除または並べ替えすることにより、モデルステップのクリティカルパスを短縮、または他の最適化を有効化します。 Identity などの実質的な NoOp ノードも削除します。
  • プルーニングオプティマイザ - グラフからの出力に影響を与えないノードを削除します。通常はこれを最初に実行することにより、グラフのサイズを縮小化し、他の Grappler パスの処理を高速化します。
  • 関数オプティマイザ - TensorFlow プログラムの関数ライブラリを最適化し、関数本体をインライン展開することにより、他のプロシージャ間最適化を可能にします。
  • 形状オプティマイザ - 形状および形状関連の情報に関する演算のサブグラフを最適化します。
  • 自動並列オプティマイザ - バッチの次元に沿って分割することにより、グラフを自動的に並列化します。このオプティマイザはデフォルトではオフです。
  • ループオプティマイザ - ループ内からループ不変のサブグラフを引き上げ、冗長なスタック演算をループ外に移動させることにより、グラフ制御フローを最適化します。また、静的に既知のトリップカウントを使用してループを最適化することにより、静的に既知の条件付き分岐のデッドブランチを削除します。
  • スコープアロケーターオプティマイザ - スコープアロケーターを導入することにより、データの移動を削減し、一部の演算を統合します。
  • ホスト固定オプティマイザ - 小さな演算を CPU にスワップします。このオプティマイザはデフォルトではオフです。
  • 自動混合精度オプティマイザ - パフォーマンスを向上させるために、適用可能な場合にデータ型を float16 に変換します。現在、GPU にのみ適用されます。
  • デバッグストリッパー - tf.debugging.Asserttf.debugging.check_numericstf.printなどのデバッグ演算関連のノードをストリップします。このオプティマイザはデフォルトではオフです。

セットアップ

import numpy as np
import timeit
import traceback
import contextlib


import tensorflow as tf

オプティマイザの状態を簡単に切り替えるためのコンテキストマネージャを作成します。

@contextlib.contextmanager
def options(options):
  old_opts = tf.config.optimizer.get_experimental_options()
  tf.config.optimizer.set_experimental_options(options)
  try:
    yield
  finally:
    tf.config.optimizer.set_experimental_options(old_opts)

Grappler 使用の有無による実行パフォーマンスを比較する

TensorFlow 2 以降は、デフォルトで Eager Execution が実装されています。デフォルトの実行を Graph モードに切り替えるには、デフォルトでtf.functionを使用します。Grappler はバックグラウンドで自動的に実行し、上記のグラフオプティマイザを適用して実行パフォーマンスを向上させます。

定数折り畳みオプティマイザ

導入的な例として、定数に対して演算を実行し、出力を返す関数を 1 つ考えます。

def test_function_1():
  @tf.function
  def simple_function(input_arg):
    print('Tracing!')
    a = tf.constant(np.random.randn(2000,2000), dtype = tf.float32)
    c = a
    for n in range(50):
      c = c@a
    return tf.reduce_mean(c+input_arg)

  return simple_function

定数折り畳みオプティマイザをオフにして、関数を実行します。

with options({'constant_folding': False}):
  print(tf.config.optimizer.get_experimental_options())
  simple_function = test_function_1()
  # Trace once
  x = tf.constant(2.2)
  simple_function(x)
  print("Vanilla execution:", timeit.timeit(lambda: simple_function(x), number = 1), "s")
{'constant_folding': False, 'disable_model_pruning': False, 'disable_meta_optimizer': False}
Tracing!
Vanilla execution: 0.06577762899996742 s

定数折り畳みオプティマイザを有効にして関数を再度実行し、関数実行の高速化を観察します。

with options({'constant_folding': True}):
  print(tf.config.optimizer.get_experimental_options())
  simple_function = test_function_1()
  # Trace once
  x = tf.constant(2.2)
  simple_function(x)
  print("Constant folded execution:", timeit.timeit(lambda: simple_function(x), number = 1), "s")
{'constant_folding': True, 'disable_model_pruning': False, 'disable_meta_optimizer': False}
Tracing!
Constant folded execution: 0.0008331830003953655 s

デバッグ ストリッパー オプティマイザ

入力引数の数値をチェックし、それを返す単純な関数を見てみましょう。

def test_function_2():
  @tf.function
  def simple_func(input_arg):
    output = input_arg
    tf.debugging.check_numerics(output, "Bad!")
    return output
  return simple_func

まず最初に、デバッグストリッパーオプティマイザをオフにして関数を実行します。

test_func = test_function_2()
p1 = tf.constant(float('inf'))
try:
  test_func(p1)
except tf.errors.InvalidArgumentError as e:
  traceback.print_exc(limit=2)
Traceback (most recent call last):
  File "<ipython-input-1-1ac473fdfbab>", line 4, in <module>
    test_func(p1)
  File "/tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py", line 828, in __call__
    result = self._call(*args, **kwds)
tensorflow.python.framework.errors_impl.InvalidArgumentError: 2 root error(s) found.
  (0) Invalid argument:  Bad! : Tensor had Inf values
     [[node CheckNumerics (defined at <ipython-input-1-cbee1561c83e>:5) ]]
     [[Identity/_4]]
  (1) Invalid argument:  Bad! : Tensor had Inf values
     [[node CheckNumerics (defined at <ipython-input-1-cbee1561c83e>:5) ]]
0 successful operations.
0 derived errors ignored. [Op:__inference_simple_func_131]

Errors may have originated from an input operation.
Input Source operations connected to node CheckNumerics:
 input_arg (defined at <ipython-input-1-1ac473fdfbab>:4)

Input Source operations connected to node CheckNumerics:
 input_arg (defined at <ipython-input-1-1ac473fdfbab>:4)

Function call stack:
simple_func -> simple_func


test_funcに対するInf引数のため、tf.debugging.check_numericsには無効な引数エラーが発生します。

デバッグストリッパーオプティマイザを有効にして、関数を再度実行します。

with options({'debug_stripper': True}):
  test_func2 = test_function_2()
  p1 = tf.constant(float('inf'))
  try:
    test_func2(p1)
  except tf.errors.InvalidArgumentError as e:
    traceback.print_exc(limit=2)

デバッグストリッパーオプティマイザがグラフからtf.debug.check_numericsノードを取り除き、エラーを発生させることなく関数を実行します。

まとめ

TensorFlow ランタイムは Grappler を使用して、実行前にグラフを自動的に最適化します。tf.config.optimizer.set_experimental_optionsを使用すると、様々なグラフオプティマイザを有効または無効にすることができます。

Grappler に関する詳しい情報は、TensorFlow グラフの最適化をご覧ください。