TensorFlow 图优化与 Grappler

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看源代码

下载笔记本

概述

TensorFlow 使用图执行和急切执行来执行计算。一个 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 及更高版本默认情况下以急切模式执行。使用 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")

启用常量折叠优化器并再次执行函数以观察函数执行速度的提高。

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")

调试剥离器优化器

考虑一个简单的函数，它检查其输入参数的数值并返回它。

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)

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