使用 tf.function 提升性能

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

在 TensorFlow 2 中，Eager Execution 默认处于启用状态。界面非常灵活直观（执行一次性运算要简单快速得多），不过，这可能对性能和可部署性造成一定影响。

您可以使用 tf.function 将程序转换为计算图。这是一个转换工具，用于从 Python 代码创建独立于 Python 的数据流图。它可以帮助您创建高效且可移植的模型，并且如果要使用 SavedModel，则必须使用此工具。

本指南介绍 tf.function 的底层工作原理，让您形成概念化理解，从而有效地加以利用。

要点和建议包括：

  • 先在 Eager 模式下调试，然后使用 @tf.function 进行装饰。
  • 不依赖 Python 的副作用，如对象变异或列表追加。
  • tf.function 最适合处理 TensorFlow 运算；NumPy 和 Python 调用会转换为常量。

设置

# Update TensorFlow, as this notebook requires version 2.9 or later
!pip install -q -U tensorflow>=2.9.0
import tensorflow as tf

2022-12-14 22:33:48.348405: 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 22:33:48.348501: 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 22:33:48.348510: 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.

定义一个辅助函数来演示可能遇到的错误类型：

import traceback
import contextlib

# Some helper code to demonstrate the kinds of errors you might encounter.
@contextlib.contextmanager
def assert_raises(error_class):
  try:
    yield
  except error_class as e:
    print('Caught expected exception \n  {}:'.format(error_class))
    traceback.print_exc(limit=2)
  except Exception as e:
    raise e
  else:
    raise Exception('Expected {} to be raised but no error was raised!'.format(
        error_class))

基础知识

用法

您定义的 Function（例如，通过应用 @tf.function 装饰器）就像核心 TensorFlow 运算：您可以在 Eager 模式下执行它，可以计算梯度，等等。

@tf.function  # The decorator converts `add` into a `Function`.
def add(a, b):
  return a + b

add(tf.ones([2, 2]), tf.ones([2, 2]))  #  [[2., 2.], [2., 2.]]

<tf.Tensor: shape=(2, 2), dtype=float32, numpy=
array([[2., 2.],
       [2., 2.]], dtype=float32)>

v = tf.Variable(1.0)
with tf.GradientTape() as tape:
  result = add(v, 1.0)
tape.gradient(result, v)

<tf.Tensor: shape=(), dtype=float32, numpy=1.0>

Function 中可以嵌套其他 Function。

@tf.function
def dense_layer(x, w, b):
  return add(tf.matmul(x, w), b)

dense_layer(tf.ones([3, 2]), tf.ones([2, 2]), tf.ones([2]))

<tf.Tensor: shape=(3, 2), dtype=float32, numpy=
array([[3., 3.],
       [3., 3.],
       [3., 3.]], dtype=float32)>

Function 的执行速度比 Eager 代码快，尤其是对于包含很多简单运算的计算图。但是，对于包含一些复杂运算（如卷积）的计算图，速度提升不会太明显。

import timeit
conv_layer = tf.keras.layers.Conv2D(100, 3)

@tf.function
def conv_fn(image):
  return conv_layer(image)

image = tf.zeros([1, 200, 200, 100])
# Warm up
conv_layer(image); conv_fn(image)
print("Eager conv:", timeit.timeit(lambda: conv_layer(image), number=10))
print("Function conv:", timeit.timeit(lambda: conv_fn(image), number=10))
print("Note how there's not much difference in performance for convolutions")

Eager conv: 0.006609583000681596
Function conv: 0.006563433998962864
Note how there's not much difference in performance for convolutions

跟踪

本部分介绍了 Function 的幕后运作方式，包括未来可能会发生变化的实现细节。但是，当您了解跟踪的原因和时间后，就能够更轻松高效地使用 tf.function！

什么是“跟踪”？

Function 在 TensorFlow 计算图中运行您的程序。但是，tf.Graph 不能代表您在 Eager TensorFlow 程序中编写的全部内容。例如，Python 支持多态，但是 tf.Graph 要求其输入具有指定的数据类型和维度。或者，您可能执行辅助任务，例如读取命令行参数、引发错误或使用更复杂的 Python 对象。这些内容均不能在 tf.Graph 中运行。

Function 通过将代码分为以下两个阶段填补了这一空缺：

  1. 第一阶段称为跟踪，在这一阶段中，Function 会创建新的 tf.Graph。Python 代码可以正常运行，但是所有 TensorFlow 运算（例如添加两个张量）都会被推迟：它们会被 tf.Graph 捕获而不运行。

  2. 在第二阶段中，将运行包含第一阶段中推迟的全部内容的 tf.Graph。此阶段比跟踪阶段快得多。

根据输入，Function 在调用时并非总会运行第一阶段。请参阅下方的跟踪规则以更好地了解其决定方式。跳过第一阶段并仅执行第二阶段，可以实现 TensorFlow 的高性能。

当 Function 决定跟踪时，在跟踪阶段完成后会立即运行第二阶段，因此调用 Function 会创建并运行 tf.Graph。稍后，您将了解如何使用 get_concrete_function 来仅运行跟踪阶段。

当您将不同类型的参数传递给 Function 时，两个阶段都将运行：

@tf.function
def double(a):
  print("Tracing with", a)
  return a + a

print(double(tf.constant(1)))
print()
print(double(tf.constant(1.1)))
print()
print(double(tf.constant("a")))
print()

Tracing with Tensor("a:0", shape=(), dtype=int32)
tf.Tensor(2, shape=(), dtype=int32)

Tracing with Tensor("a:0", shape=(), dtype=float32)
tf.Tensor(2.2, shape=(), dtype=float32)

Tracing with Tensor("a:0", shape=(), dtype=string)
tf.Tensor(b'aa', shape=(), dtype=string)

请注意，如果重复使用同一参数类型调用 Function，TensorFlow 会跳过跟踪阶段并重用之前跟踪的计算图，因为后面的调用生成的计算图可能相同。

# This doesn't print 'Tracing with ...'
print(double(tf.constant("b")))

tf.Tensor(b'bb', shape=(), dtype=string)

您可以使用 pretty_printed_concrete_signatures() 查看所有可用跟踪记录：

print(double.pretty_printed_concrete_signatures())

double(a)
  Args:
    a: int32 Tensor, shape=()
  Returns:
    int32 Tensor, shape=()

double(a)
  Args:
    a: float32 Tensor, shape=()
  Returns:
    float32 Tensor, shape=()

double(a)
  Args:
    a: string Tensor, shape=()
  Returns:
    string Tensor, shape=()

目前，您已经了解 tf.function 通过 TensorFlow 的计算图跟踪逻辑创建缓存的动态调度层。对于术语的含义，更具体的解释如下：

  • tf.Graph 与语言无关，是 TensorFlow 计算的原始可移植表示。
  • ConcreteFunction 封装 tf.Graph。
  • Function 管理 ConcreteFunction 的缓存，并为输入选择正确的缓存。
  • tf.function 包装 Python 函数，并返回一个 Function 对象。
  • 跟踪会创建 tf.Graph 并将其封装在 ConcreteFunction 中，也称为跟踪。

跟踪规则

被调用时，Function 使用每个参数的 tf.types.experimental.TraceType 将调用参数与现有的 ConcreteFunction 匹配。如果找到匹配的 ConcreteFunction，则将调用分派给它。如果未找到匹配项，则跟踪新的 ConcreteFunction。

如果找到多个匹配项，则会选择最具体的签名。匹配是通过子类型化完成的，就像 C++ 或 Java 中的普通函数调用一样。例如，TensorShape([1, 2]) 是 TensorShape([None, None]) 的子类型，因此可以将使用 TensorShape([1, 2]) 对 tf.function 进行的调用分派到使用 TensorShape([None, None]) 生成的 ConcreteFunction。但是，如果具有 TensorShape([1, None]) 的 ConcreteFunction 也存在，那么它将被优先考虑，因为它更具体。

TraceType 由输入参数确定，具体如下所示：

  • 对于 Tensor，类型由 Tensor 的 dtype 和 shape 参数化；有秩形状是无秩形状的子类型；固定维度是未知维度的子类型
  • 对于 Variable，类型类似于 Tensor，但还包括变量的唯一资源 ID，这是正确连接控制依赖项所必需的
  • 对于 Python 基元值，类型对应于值本身。例如，值为 3 的 TraceType 是 LiteralTraceType<3>，而不是 int。
  • 对于 list 和 tuple 等 Python 有序容器，类型是通过其元素的类型来参数化的；例如，[1, 2] 的类型是 ListTraceType<LiteralTraceType<1>, LiteralTraceType<2>>，[2, 1] 的类型是 ListTraceType<LiteralTraceType<2>, LiteralTraceType<1>>，两者不同。
  • 对于 dict 等 Python 映射，类型也是从相同的键到值类型而不是实际值的映射。例如，{1: 2, 3: 4} 的类型为 MappingTraceType<<KeyValue<1, LiteralTraceType<2>>>, <KeyValue<3, LiteralTraceType<4>>>>。但是，与有序容器不同的是，{1: 2, 3: 4} 和 {3: 4, 1: 2} 具有等价的类型。
  • 对于实现 __tf_tracing_type__ 方法的 Python 对象，类型为该方法返回的任何内容
  • 对于任何其他 Python 对象，类型是通用的 TraceType，它使用对象的 Python 相等性和散列进行匹配。（注：它依赖于对对象的弱引用，因此仅在对象处于范围内/未被删除时才有效。）

注：TraceType 基于 Function 输入参数，因此仅对全局变量和自由变量进行更改将不会创建新的跟踪记录。有关处理 Python 全局变量和自由变量的建议做法，请参阅本部分。

控制回溯

回溯即 Function 创建多个跟踪记录的过程，可以确保 TensorFlow 为每组输入生成正确的计算图。但是，跟踪非常消耗资源！如果 Function 为每一次调用都回溯新的计算图，您会发现代码的执行速度远不如不使用 tf.function 时快。

要控制跟踪行为，可以采用以下技巧：

将固定的 input_signature 传递给 tf.function

@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))
def next_collatz(x):
  print("Tracing with", x)
  return tf.where(x % 2 == 0, x // 2, 3 * x + 1)

print(next_collatz(tf.constant([1, 2])))
# You specified a 1-D tensor in the input signature, so this should fail.
with assert_raises(ValueError):
  next_collatz(tf.constant([[1, 2], [3, 4]]))

# You specified an int32 dtype in the input signature, so this should fail.
with assert_raises(ValueError):
  next_collatz(tf.constant([