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

图和函数简介

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

图和tf.function简介

本指南深入研究TensorFlow和Keras的表面,以了解TensorFlow的工作原理。如果您想立即开始使用Keras,请参阅我们的Keras指南集合

在本指南中,您将看到TensorFlow如何允许您对代码进行简单更改以获取图形,如何存储和表示图形以及如何使用它们来加速和导出模型的核心。

这是一个简短的介绍;有关这些概念的完整介绍,请参见tf.function指南

什么是图?

在之前的三本指南中,您已经看到TensorFlow 急切地运行。这意味着TensorFlow操作由Python执行,然后逐个操作,然后将结果返回给Python。 Eager TensorFlow利用了GPU的优势,允许您在GPU和TPU上放置变量,张量甚至操作。调试也很容易。

对于某些用户,您可能永远不需要或不想离开Python。

但是,在Python中逐个运行TensorFlow会阻止大量其他可用的加速。如果可以从Python提取张量计算,则可以将它们制成

图是tf.Operation对象的数据结构, tf.Operation对象表示计算单位; tf.Tensor对象,代表操作之间流动的数据单位。它们在tf.Graph上下文中定义。由于这些图是数据结构,因此可以保存,运行和还原它们,而无需原始的Python代码。

在TensorBoard中可视化时,这就是一个简单的两层图。

两层张量流图

图的好处

使用图形,您将拥有很大的灵活性。您可以在没有Python解释器的环境(如移动应用程序,嵌入式设备和后端服务器)中使用TensorFlow图。 TensorFlow从Python导出图形时,会将图形用作已保存模型的格式。

图形也很容易优化,从而允许编译器进行如下转换:

  • 通过在计算中折叠常数节点来静态推断张量的值(“常数折叠”)
  • 独立的计算子部分,在线程或设备之间进行拆分。
  • 通过消除常见的子表达式来简化算术运算。

有一个完整的优化系统Grappler来执行此和其他加速操作。

简而言之,图形非常有用,它可以使TensorFlow 快速运行,并行运行并在多个设备上高效运行。

但是,为了方便起见,您仍然想在Python中定义我们的机器学习模型(或其他计算),然后在需要时自动构造图。

追踪图

在TensorFlow中创建图形的方式是将tf.function用作直接调用或装饰器。

 import tensorflow as tf
import timeit
from datetime import datetime
 
 # Define a Python function
def function_to_get_faster(x, y, b):
  x = tf.matmul(x, y)
  x = x + b
  return x

# Create a `Function` object that contains a graph
a_function_that_uses_a_graph = tf.function(function_to_get_faster)

# Make some tensors
x1 = tf.constant([[1.0, 2.0]])
y1 = tf.constant([[2.0], [3.0]])
b1 = tf.constant(4.0)

# It just works!
a_function_that_uses_a_graph(x1, y1, b1).numpy()
 
array([[12.]], dtype=float32)

tf.function -ized函数是Python可调用的函数,其功能与Python等效项相同。它们具有特定的类( python.eager.def_function.Function ),但对您而言,它们仅充当非跟踪版本。

tf.function递归地跟踪它调用的任何Python函数。

 def inner_function(x, y, b):
  x = tf.matmul(x, y)
  x = x + b
  return x

# Use the decorator
@tf.function
def outer_function(x):
  y = tf.constant([[2.0], [3.0]])
  b = tf.constant(4.0)

  return inner_function(x, y, b)

# Note that the callable will create a graph that
# includes inner_function() as well as outer_function()
outer_function(tf.constant([[1.0, 2.0]])).numpy()
 
array([[12.]], dtype=float32)

如果您使用过TensorFlow 1.x,您会发现您根本不需要定义一个Placeholdertf.Sesssion

流量控制和副作用

默认情况下,流控制和循环通过tf.autograph转换为TensorFlow。 Autograph使用了多种方法的组合,包括标准化循环结构,展开和AST操作。

 def my_function(x):
  if tf.reduce_sum(x) <= 1:
    return x * x
  else:
    return x-1

a_function = tf.function(my_function)

print("First branch, with graph:", a_function(tf.constant(1.0)).numpy())
print("Second branch, with graph:", a_function(tf.constant([5.0, 5.0])).numpy())
 
First branch, with graph: 1.0
Second branch, with graph: [4. 4.]

您可以直接调用Autograph转换以查看Python如何转换为TensorFlow ops。大多数情况下,这是不可读的,但是您可以看到转换。

 # Don't read the output too carefully.
print(tf.autograph.to_code(my_function))
 
def tf__my_function(x):
    with ag__.FunctionScope('my_function', 'fscope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as fscope:
        do_return = False
        retval_ = ag__.UndefinedReturnValue()

        def get_state():
            return (retval_, do_return)

        def set_state(vars_):
            nonlocal retval_, do_return
            (retval_, do_return) = vars_

        def if_body():
            nonlocal retval_, do_return
            try:
                do_return = True
                retval_ = (ag__.ld(x) * ag__.ld(x))
            except:
                do_return = False
                raise

        def else_body():
            nonlocal retval_, do_return
            try:
                do_return = True
                retval_ = (ag__.ld(x) - 1)
            except:
                do_return = False
                raise
        ag__.if_stmt((ag__.converted_call(ag__.ld(tf).reduce_sum, (ag__.ld(x),), None, fscope) <= 1), if_body, else_body, get_state, set_state, ('retval_', 'do_return'), 2)
        return fscope.ret(retval_, do_return)


Autograph自动转换if-then子句,循环, breakreturncontinue等。

在大多数情况下,Autograph无需任何特殊考虑即可工作。但是,有一些警告, tf.function指南可以在这里提供帮助,以及完整的签名参考

看到速度

仅将使用张量的函数包装在tf.function并不会自动加速代码。对于在单个计算机上多次调用的小函数,调用图形或图形片段的开销可能会占主导地位。同样,如果大多数计算已经在加速器上进行,例如大量GPU的卷积,则图形加速不会很大。

对于复杂的计算,图形可以显着提高速度。这是因为图形减少了Python与设备之间的通信,并加快了速度。

此代码在一些小的密集层上运行几次。

 # Create an oveerride model to classify pictures
class SequentialModel(tf.keras.Model):
  def __init__(self, **kwargs):
    super(SequentialModel, self).__init__(**kwargs)
    self.flatten = tf.keras.layers.Flatten(input_shape=(28, 28))
    self.dense_1 = tf.keras.layers.Dense(128, activation="relu")
    self.dropout = tf.keras.layers.Dropout(0.2)
    self.dense_2 = tf.keras.layers.Dense(10)

  def call(self, x):
    x = self.flatten(x)
    x = self.dense_1(x)
    x = self.dropout(x)
    x = self.dense_2(x)
    return x

input_data = tf.random.uniform([60, 28, 28])

eager_model = SequentialModel()
graph_model = tf.function(eager_model)

print("Eager time:", timeit.timeit(lambda: eager_model(input_data), number=10000))
print("Graph time:", timeit.timeit(lambda: graph_model(input_data), number=10000))

 
Eager time: 4.9249971129993355
Graph time: 2.026840765000088

多态函数

跟踪函数时,创建的Function对象是polymorphic 。多态函数是可调用的Python,它在一个API后面封装了几个具体的函数图。

您可以对所有不同类型的dtypes和shape使用此Function 。每次使用新的参数签名调用它时,原始函数都会使用新的参数重新跟踪。然后, Function将与该tf.Graph相对应的tf.Graph存储在concrete_function 。如果已经使用该类型的参数跟踪了函数,则只需获取预跟踪图即可。

然后,从概念上讲:

  • tf.Graph是描述计算的原始可移植数据结构
  • Function是ConcreteFunctions上的缓存,跟踪,调度程序
  • ConcreteFunction是围绕图的热切兼容包装,可让您从Python执行图

检查多态函数

您可以检查a_function ,这是在Python函数my_function上调用tf.function的结果。在此示例中,使用三种参数调用a_function导致三个不同的具体函数。

 print(a_function)

print("Calling a `Function`:")
print("Int:", a_function(tf.constant(2)))
print("Float:", a_function(tf.constant(2.0)))
print("Rank-1 tensor of floats", a_function(tf.constant([2.0, 2.0, 2.0])))
 
<tensorflow.python.eager.def_function.Function object at 0x7f466417bf60>
Calling a `Function`:
Int: tf.Tensor(1, shape=(), dtype=int32)
Float: tf.Tensor(1.0, shape=(), dtype=float32)
Rank-1 tensor of floats tf.Tensor([1. 1. 1.], shape=(3,), dtype=float32)

 # Get the concrete function that works on floats
print("Inspecting concrete functions")
print("Concrete function for float:")
print(a_function.get_concrete_function(tf.TensorSpec(shape=[], dtype=tf.float32)))
print("Concrete function for tensor of floats:")
print(a_function.get_concrete_function(tf.constant([2.0, 2.0, 2.0])))

 
Inspecting concrete functions
Concrete function for float:
ConcreteFunction my_function(x)
  Args:
    x: float32 Tensor, shape=()
  Returns:
    float32 Tensor, shape=()
Concrete function for tensor of floats:
ConcreteFunction my_function(x)
  Args:
    x: float32 Tensor, shape=(3,)
  Returns:
    float32 Tensor, shape=(3,)

 # Concrete functions are callable
# Note: You won't normally do this, but instead just call the containing `Function`
cf = a_function.get_concrete_function(tf.constant(2))
print("Directly calling a concrete function:", cf(tf.constant(2)))
 
Directly calling a concrete function: tf.Tensor(1, shape=(), dtype=int32)

在此示例中,您看到的内容非常丰富。除非您专门管理跟踪,否则通常将不需要直接调用具体函数,如此处所示。

恢复渴望执行

您可能会发现自己正在查看长堆栈跟踪,特别是那些引用tf.Graphwith tf.Graph().as_default() 。这意味着您很可能在图形上下文中运行。 TensorFlow中的核心功能使用图上下文,例如model.fit()model.fit()

调试急切的执行通常要容易得多。堆栈跟踪应相对较短,并且易于理解。

在图形使调试变得棘手的情况下,您可以恢复使用急切的执行进行调试。

您可以通过以下方法确保自己热切运行:

  • 直接将模型和图层称为可调用对象

  • 使用Keras编译/拟合时,在编译时使用model.compile(run_eagerly=True)

  • 通过tf.config.run_functions_eagerly(True)设置全局执行模式

使用run_eagerly=True

 # Define an identity layer with an eager side effect
class EagerLayer(tf.keras.layers.Layer):
  def __init__(self, **kwargs):
    super(EagerLayer, self).__init__(**kwargs)
    # Do some kind of initialization here

  def call(self, inputs):
    print("\nCurrently running eagerly", str(datetime.now()))
    return inputs
 
 # Create an override model to classify pictures, adding the custom layer
class SequentialModel(tf.keras.Model):
  def __init__(self):
    super(SequentialModel, self).__init__()
    self.flatten = tf.keras.layers.Flatten(input_shape=(28, 28))
    self.dense_1 = tf.keras.layers.Dense(128, activation="relu")
    self.dropout = tf.keras.layers.Dropout(0.2)
    self.dense_2 = tf.keras.layers.Dense(10)
    self.eager = EagerLayer()

  def call(self, x):
    x = self.flatten(x)
    x = self.dense_1(x)
    x = self.dropout(x)
    x = self.dense_2(x)
    return self.eager(x)

# Create an instance of this model
model = SequentialModel()

# Generate some nonsense pictures and labels
input_data = tf.random.uniform([60, 28, 28])
labels = tf.random.uniform([60])

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
 

首先,不要急于编译模型。请注意,没有追踪模型;尽管有其名称, compile仅设置损失函数,优化和其他训练参数。

 model.compile(run_eagerly=False, loss=loss_fn)
 

现在,调用fit并查看该函数已被跟踪(两次),然后急切的效果不再运行。

 model.fit(input_data, labels, epochs=3)
 
Epoch 1/3

Currently running eagerly 2020-08-04 01:22:21.848492

Currently running eagerly 2020-08-04 01:22:21.955102
2/2 [==============================] - 0s 1ms/step - loss: 1.4056
Epoch 2/3
2/2 [==============================] - 0s 1ms/step - loss: 0.0037
Epoch 3/3
2/2 [==============================] - 0s 1ms/step - loss: 0.0019

<tensorflow.python.keras.callbacks.History at 0x7f45f03c5630>

但是,如果您只渴望一个时期,您就会看到两次渴望的副作用。

 print("Running eagerly")
# When compiling the model, set it to run eagerly
model.compile(run_eagerly=True, loss=loss_fn)

model.fit(input_data, labels, epochs=1)

 
Running eagerly

Currently running eagerly 2020-08-04 01:22:22.152979
1/2 [==============>...............] - ETA: 0s - loss: 8.7806e-04
Currently running eagerly 2020-08-04 01:22:22.173295
2/2 [==============================] - 0s 5ms/step - loss: 4.6877e-04

<tensorflow.python.keras.callbacks.History at 0x7f45f0312940>

使用run_functions_eagerly

您还可以全局设置所有内容以使其热切运行。请注意,这仅在您重新跟踪时才有效;跟踪的函数将保持跟踪状态并作为图形运行。

 # Now, globally set everything to run eagerly
tf.config.run_functions_eagerly(True)
print("Run all functions eagerly.")

# First, trace the model, triggering the side effect
polymorphic_function = tf.function(model)

# It was traced...
print(polymorphic_function.get_concrete_function(input_data))

# But when you run the function again, the side effect happens (both times).
result = polymorphic_function(input_data)
result = polymorphic_function(input_data)
 
Run all functions eagerly.

Currently running eagerly 2020-08-04 01:22:22.202726
ConcreteFunction function(self)
  Args:
    self: float32 Tensor, shape=(60, 28, 28)
  Returns:
    float32 Tensor, shape=(60, 10)

Currently running eagerly 2020-08-04 01:22:22.206521

Currently running eagerly 2020-08-04 01:22:22.207818

 # Don't forget to set it back when you are done
tf.config.experimental_run_functions_eagerly(False)

 
WARNING:tensorflow:From <ipython-input-17-782fe9ce7b18>:2: experimental_run_functions_eagerly (from tensorflow.python.eager.def_function) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.config.run_functions_eagerly` instead of the experimental version.

追踪与表现

跟踪会花费一些开销。尽管快速跟踪小型功能,但是大型模型可能需要花费大量的时间来跟踪。通常可以通过提高性能来迅速收回投资,但是必须意识到,由于跟踪,任何大型模型训练的前几个时期可能会变慢。

无论模型多大,都希望避免频繁跟踪。 tf.function指南的这一部分讨论了如何设置输入规范并使用张量参数来避免跟踪。如果发现性能异常变差,最好检查一下是否意外还原。

您可以添加仅渴望的副作用(例如打印Python参数),以便查看何时跟踪该函数。在这里,您会看到额外的跟踪,因为新的Python参数始终会触发跟踪。

 # Use @tf.function decorator
@tf.function
def a_function_with_python_side_effect(x):
  print("Tracing!")  # This eager
  return x * x + tf.constant(2)

# This is traced the first time
print(a_function_with_python_side_effect(tf.constant(2)))
# The second time through, you won't see the side effect
print(a_function_with_python_side_effect(tf.constant(3)))

# This retraces each time the Python argument chances
# as a Python argument could be an epoch count or other
# hyperparameter
print(a_function_with_python_side_effect(2))
print(a_function_with_python_side_effect(3))

 
Tracing!
tf.Tensor(6, shape=(), dtype=int32)
tf.Tensor(11, shape=(), dtype=int32)
Tracing!
tf.Tensor(6, shape=(), dtype=int32)
Tracing!
tf.Tensor(11, shape=(), dtype=int32)

下一步

您可以在tf.function API参考页和指南上阅读更深入的讨论。