使用 Grappler 优化 TensorFlow 计算图

##### Copyright 2020 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

使用 Grappler 优化 TensorFlow 计算图#

在 TensorFlow.org 上查看 在 Google Colab 中运行 在 Github 上查看源代码 下载笔记本 
%cd ..
from set_env import temp_dir

/media/pc/data/lxw/ai/d2py/doc/libs/tf-chaos/guide

概述#

TensorFlow 同时使用计算图和 Eager Execution 来执行计算。一个 tf.Graph 包含一组代表计算单元的 tf.Operation 对象（运算）和一组代表在运算之间流动的数据单元的 tf.Tensor 对象。

Grappler 是 TensorFlow 运行时中的默认计算图优化系统。Grappler 通过计算图简化和其他高级优化（例如利用内嵌函数体实现程序间优化），在计算图模式（在 tf.function 内）下应用优化以提高 TensorFlow 计算的性能。优化 tf.Graph 还可以通过优化计算图节点到计算资源的映射来减少设备峰值内存使用量并提高硬件利用率。

使用 tf.config.optimizer.set_experimental_options() 可以更好地控制 tf.Graph 优化。

可用的计算图优化器#

Grappler 通过称为 MetaOptimizer 的顶级驱动程序执行计算图优化。TensorFlow 提供以下计算图优化器：

  • 常量折叠优化器 - 通过折叠计算图中的常量节点来静态推断张量的值（如可能），并使用常量使结果具体化。

  • 算术优化器 - 通过消除常见的子表达式并简化算术语句来简化算术运算。

  • 布局优化器 - 优化张量布局以更高效地执行依赖于数据格式的运算，例如卷积。

  • 重新映射优化器 - 通过将常见的子计算图替换为经过优化的融合一体化内核，将子计算图重新映射到更高效的实现上。

  • 内存优化器 - 分析计算图以检查每个运算的峰值内存使用量，并插入 CPU-GPU 内存复制操作以将 GPU 内存交换到 CPU，从而减少峰值内存使用量。

  • 依赖项优化器 - 移除或重新排列控制依赖项，以缩短模型步骤的关键路径或实现其他优化。另外，还移除了实际上是无运算的节点，例如 Identity。

  • 剪枝优化器 - 修剪对计算图的输出没有影响的节点。通常会首先运行剪枝来减小计算图的大小并加快其他 Grappler 传递中的处理速度。

  • 函数优化器 - 优化 TensorFlow 程序的函数库，并内嵌函数体以实现其他程序间优化。

  • 形状优化器 - 优化对形状和形状相关信息进行运算的子计算图。

  • 自动并行优化器 - 通过沿批次维度拆分来自动并行化计算图。默认情况下，此优化器处于关闭状态。

  • 循环优化器 - 通过将循环不变式子计算图提升到循环外并通过移除循环中的冗余堆栈运算来优化计算图控制流。另外，还优化具有静态已知行程计数的循环，并移除条件语句中静态已知的无效分支。

  • 范围分配器优化器 - 引入范围分配器以减少数据移动并合并某些运算。

  • 固定到主机优化器 - 将小型运算交换到 CPU 上。默认情况下，此优化器处于关闭状态。

  • 自动混合精度优化器 - 在适用的情况下将数据类型转换为 float16 以提高性能。目前仅适用于 GPU。

  • 调试剥离器 - 从计算图中剥离与调试运算相关的节点，例如 tf.debugging.Assert、tf.debugging.check_numerics 和 tf.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 模式执行。使用 tf.function 可将默认执行切换为“计算图”模式。Grappler 在后台自动运行，以应用上述计算图优化并提高执行性能。

常量折叠优化器#

作为一个初步的示例，考虑一个对常量执行运算并返回输出的函数。

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.0024620250333100557 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.0006806210149079561 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 "/tmp/ipykernel_4128893/3616845043.py", line 4, in <module>
    test_func(p1)
  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/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 "<frozen runpy>", line 198, in _run_module_as_main

  File "<frozen runpy>", line 88, in _run_code

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/ipykernel_launcher.py", line 17, in <module>

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/traitlets/config/application.py", line 1075, in launch_instance

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/ipykernel/kernelapp.py", line 701, in start

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/tornado/platform/asyncio.py", line 205, in start

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/asyncio/base_events.py", line 639, in run_forever

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/asyncio/base_events.py", line 1985, in _run_once

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/asyncio/events.py", line 88, in _run

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 534, in dispatch_queue

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 523, in process_one

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 429, in dispatch_shell

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 767, in execute_request

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/ipykernel/ipkernel.py", line 429, in do_execute

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/ipykernel/zmqshell.py", line 549, in run_cell

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3075, in run_cell

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3130, in _run_cell

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/IPython/core/async_helpers.py", line 128, in _pseudo_sync_runner

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3334, in run_cell_async

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3517, in run_ast_nodes

  File "/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3577, in run_code

  File "/tmp/ipykernel_4128893/3616845043.py", line 4, in <module>

  File "/tmp/ipykernel_4128893/2241890286.py", line 5, in simple_func

Bad! : Tensor had Inf values
	 [[{{node CheckNumerics}}]] [Op:__inference_simple_func_128]

由于 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 计算图优化。