此页面由 Cloud Translation API 翻译。
Switch to English

使用Grappler优化TensorFlow图

在TensorFlow.org上查看 在Google Colab中运行 在GitHub上查看源代码 下载笔记本

总览

TensorFlow同时使用图形执行和急切执行来执行计算。一个tf.Graph包含一组代表计算单位的tf.Operation对象(ops)和一tf.Tensor代表在ops之间流动的数据单位的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,以减少峰值内存使用情况。
  • 依赖优化器-删除或重新排列控件依赖关系,以缩短模型步骤的关键路径或启用其他优化。还删除实际上是无操作的节点,例如身份。
  • 修剪优化器-修剪对图形输出没有影响的节点。通常首先运行它以减小图形的大小并加快其他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及更高版本默认情况下会急切执行。使用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.004905937999865273 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.0005699149999145448 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-8-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 580, 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-7-cbee1561c83e>:5) ]]
     [[Identity/_4]]
  (1) Invalid argument:  Bad! : Tensor had Inf values
     [[node CheckNumerics (defined at <ipython-input-7-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-8-1ac473fdfbab>:4)

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

Function call stack:
simple_func -> simple_func


tf.debugging.check_numerics提高,因为一个无效的参数错误Inf参数test_func

启用调试剥离程序优化器,然后再次执行该功能。

 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图优化