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
2022-12-14 20:49:48.539011: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2022-12-14 20:49:48.539107: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2022-12-14 20:49:48.539117: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.

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

@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.0013519179997274478 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.0007739210000181629 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)
2022-12-14 20:50:08.397437: E tensorflow/core/kernels/check_numerics_op.cc:293] abnormal_detected_host @0x7f8d5a600100 = {0, 1} Bad!
Traceback (most recent call last):
  File "/tmpfs/tmp/ipykernel_74478/3616845043.py", line 4, in <module>
    test_func(p1)
  File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
tensorflow.python.framework.errors_impl.InvalidArgumentError: Graph execution error:

Detected at node 'CheckNumerics' defined at (most recent call last):
    File "/usr/lib/python3.9/runpy.py", line 197, in _run_module_as_main
      return _run_code(code, main_globals, None,
    File "/usr/lib/python3.9/runpy.py", line 87, in _run_code
      exec(code, run_globals)
    File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel_launcher.py", line 17, in <module>
      app.launch_new_instance()
    File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/traitlets/config/application.py", line 992, in launch_instance
      app.start()
    File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/kernelapp.py", line 711, in start
      self.io_loop.start()
    File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tornado/platform/asyncio.py", line 215, in start
      self.asyncio_loop.run_forever()
    File "/usr/lib/python3.9/asyncio/base_events.py", line 601, in run_forever
      self._run_once()
    File "/usr/lib/python3.9/asyncio/base_events.py", line 1905, in _run_once
      handle._run()
    File "/usr/lib/python3.9/asyncio/events.py", line 80, in _run
      self._context.run(self._callback, *self._args)
    File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/kernelbase.py", line 510, in dispatch_queue
      await self.process_one()
    File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/kernelbase.py", line 499, in process_one
      await dispatch(*args)
    File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/kernelbase.py", line 406, in dispatch_shell
      await result
    File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/kernelbase.py", line 729, in execute_request
      reply_content = await reply_content
    File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/ipkernel.py", line 411, in do_execute
      res = shell.run_cell(
    File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/zmqshell.py", line 531, in run_cell
      return super().run_cell(*args, **kwargs)
    File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 2940, in run_cell
      result = self._run_cell(
    File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 2995, in _run_cell
      return runner(coro)
    File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/IPython/core/async_helpers.py", line 129, in _pseudo_sync_runner
      coro.send(None)
    File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3194, in run_cell_async
      has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
    File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3373, in run_ast_nodes
      if await self.run_code(code, result, async_=asy):
    File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3433, in run_code
      exec(code_obj, self.user_global_ns, self.user_ns)
    File "/tmpfs/tmp/ipykernel_74478/3616845043.py", line 4, in <module>
      test_func(p1)
    File "/tmpfs/tmp/ipykernel_74478/2241890286.py", line 5, in simple_func
      tf.debugging.check_numerics(output, "Bad!")
Node: 'CheckNumerics'
Bad! : Tensor had Inf values
     [[{ {node CheckNumerics} }]] [Op:__inference_simple_func_131]

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 グラフの最適化をご覧ください。