在TensorFlow 2.0中，默认情况下会打开eager execution，这为您提供了一个非常直观和灵活的用户界面（运行一次性操作更容易、更快）但这可能会牺牲性能和可部署性。
为了获得最佳性能并使您的模型可以在任何地方部署，我们提供了tf.function作为您可以用来从程序中生成图形的工具。
from __future__ import absolute_import, division, print_function, unicode_literals
import tensorflow as tf # pip install -q tensorflow==2.0.0-alpha0
# 功能就像一个操作
@tf.function
def add(a, b):
return a + b
print(add(tf.ones([2, 2]), tf.ones([2, 2]))) # [[2., 2.], [2., 2.]]<tf.Tensor: id=16, shape=(2, 2), dtype=float32, numpy= array([[2., 2.], [2., 2.]], dtype=float32)>
你定义的tf.function就像一个核心TensorFlow操作：你可以急切地执行它、你可以在图形中使用它，它有渐变（gradient）等。
# Functions have gradients
@tf.function
def add(a, b):
return a + b
v = tf.Variable(1.0)
with tf.GradientTape() as tape:
result = add(v, 1.0)
print(tape.gradient(result, v))<tf.Tensor: id=44, shape=(), dtype=float32, numpy=1.0>
# 你可以在函数内部使用函数
@tf.function
def dense_layer(x, w, b):
return add(tf.matmul(x, w), b)
print(dense_layer(tf.ones([3, 2]), tf.ones([2, 2]), tf.ones([2])))<tf.Tensor: id=74, shape=(3, 2), dtype=float32, numpy= array([[3., 3.], [3., 3.], [3., 3.]], dtype=float32)>
多态性
tf.function试图像Python函数一样通用。您可以使用各种签名调用Python函数，Python通常会做一些合理的事情。tf.function为你做这种类型的多态，即使它生成的底层TensorFlow图特定于其签名中的特定类型。
您可以使用不同类型的参数调用函数来查看正在发生的事情。
# Functions are polymorphic
@tf.function
def add(a):
return a + a
print("add 1", add(1))
print("add 1.1", add(1.1))
print("add string tensor", add(tf.constant("a")))
c = add.get_concrete_function(tf.TensorSpec(shape=None, dtype=tf.string))
print(c(a=tf.constant("a"))) # aaadd 1 tf.Tensor(2, shape=(), dtype=int32) add 1.1 tf.Tensor(2.2, shape=(), dtype=float32) add string tensor tf.Tensor(b’aa’, shape=(), dtype=string) <tf.Tensor: id=104, shape=(), dtype=string, numpy=b’aa’>
# Functions can be faster than eager code, for graphs with many small ops
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")
lstm_cell = tf.keras.layers.LSTMCell(10)
@tf.function
def lstm_fn(input, state):
return lstm_cell(input, state)
input = tf.zeros([10, 10])
state = [tf.zeros([10, 10])] * 2
# warm up
lstm_cell(input, state)
lstm_fn(input, state)
print("eager lstm:", timeit.timeit(lambda: lstm_cell(input, state), number=10))
print("function lstm:", timeit.timeit(lambda: lstm_fn(input, state), number=10))Eager conv: 0.046443079598248005 Function conv: 0.04211280681192875 Note how there’s not much difference in performance for convolutions eager lstm: 0.030607654713094234 function lstm: 0.005289177875965834
tf.function中的状态
在一般数据流图上，作为编程模型的函数的一个非常吸引人的特性是函数可以为运行时提供有关代码的预期行为的更多信息。
例如，当编写具有多个读取和写入相同变量的代码时，数据流图可能不会自然地编码最初预期的操作顺序。但是，在tf.function中，因为我们正在转换从Python跟踪的代码，所以我们知道了预期的执行顺序。
这意味着不需要添加手动控制依赖项。tf.function非常智能，可以为代码添加最小的必要和充分的控制依赖关系，以便正确运行。
# Automatic control dependencies
a = tf.Variable(1.0)
b = tf.Variable(2.0)
@tf.function
def f(x, y):
a.assign(y * b)
b.assign_add(x * a)
return a + b
f(1.0, 2.0) # 10.0<tf.Tensor: id=1610, shape=(), dtype=float32, numpy=10.0>
变量
我们可以使用相同的想法来利用代码的预期执行顺序，以便在tf.function中很容易地创建和使用变量。然而，有一个非常重要的警告，即使用变量，可以编写在多次调用时以及多次评估其输出张量时行为不同的代码。
这是一个简单的例子：
@tf.function
def f(x):
v = tf.Variable(1.0)
v.assign_add(x)
return v
f(1.) # Note: BROKEN, will throw exception如果你以急切的执行方式运行它，你总会得到“2”作为答案，但如果你在图形上下文中反复评估从f（1.）获得的Tensor，你会得到越来越多的数字。
所以tf.function不允许你编写这样的代码。
# Non-ambiguous code is ok though
v = tf.Variable(1.0)
@tf.function
def f(x):
return v.assign_add(x)
f(1.0) # 2.0
f(2.0) # 4.0<tf.Tensor: id=1635, shape=(), dtype=float32, numpy=4.0>
# 您也可以在tf.function中创建变量，只要我们可以证明这些变量仅在第一次执行函数时创建。
class C: pass
obj = C()
obj.v = None
@tf.function
def g(x):
if obj.v is None:
obj.v = tf.Variable(1.0)
return obj.v.assign_add(x)
g(1.0) # 2.0
g(2.0) # 4.0<tf.Tensor: id=1689, shape=(), dtype=float32, numpy=4.0>
# 变量初始值设定项可以依赖于函数参数和其他变量的值。
# 我们可以使用与生成控制依赖关系相同的方法找出正确的初始化顺序。
state = []
@tf.function
def fn(x):
if not state:
state.append(tf.Variable(2.0 * x))
state.append(tf.Variable(state[0] * 3.0))
return state[0] * x * state[1]
fn(tf.constant(1.0))
fn(tf.constant(3.0))WARNING: Logging before flag parsing goes to stderr. W0411 00:29:06.727809 140132268373760 tf_logging.py:161] Entity \<method-wrapper ‘call‘ of weakref object at 0x7f71f847b458> could not be transformed and will be staged without change. Error details can be found in the logs when running with the env variable AUTOGRAPH_VERBOSITY >= 1. Please report this to the AutoGraph team. Cause: Object conversion is not yet supported. If you are trying to convert code that uses an existing object, try including the creation of that object in the conversion. For example, instead of converting the method of a class, try converting the entire class instead. See https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/README.md#using-the-functional-api for more information. W0411 00:29:06.733919 140132268373760 tf_logging.py:161] Entity \<method-wrapper ‘call‘ of weakref object at 0x7f71f847b958> could not be transformed and will be staged without change. Error details can be found in the logs when running with the env variable AUTOGRAPH_VERBOSITY >= 1. Please report this to the AutoGraph team. Cause: Object conversion is not yet supported. If you are trying to convert code that uses an existing object, try including the creation of that object in the conversion. For example, instead of converting the method of a class, try converting the entire class instead. See https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/README.md#using-the-functional-api for more information. WARNING: Entity \<method-wrapper ‘call‘ of weakref object at 0x7f71f847b458> could not be transformed and will be staged without change. Error details can be found in the logs when running with the env variable AUTOGRAPH_VERBOSITY >= 1. Please report this to the AutoGraph team. Cause: Object conversion is not yet supported. If you are trying to convert code that uses an existing object, try including the creation of that object in the conversion. For example, instead of converting the method of a class, try converting the entire class instead. See https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/README.md#using-the-functional-api for more information. WARNING: Entity \<method-wrapper ‘call‘ of weakref object at 0x7f71f847b958> could not be transformed and will be staged without change. Error details can be found in the logs when running with the env variable AUTOGRAPH_VERBOSITY >= 1. Please report this to the AutoGraph team. Cause: Object conversion is not yet supported. If you are trying to convert code that uses an existing object, try including the creation of that object in the conversion. For example, instead of converting the method of a class, try converting the entire class instead. See https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/README.md#using-the-functional-api for more information. \<tf.Tensor: id=1796, shape=(), dtype=float32, numpy=36.0>
控制流和AutoGraph
虽然tf.cond和tf.while_loop继续使用tf.function，但我们提供了一个基于Python代码的轻量级编译的更好的替代方案。
autograph库与tf.function完全集成，它将重写依赖于Tensors的条件和循环，以便在图中动态运行。
# Simple loop
@tf.function
def f(x):
while tf.reduce_sum(x) > 1:
tf.print(x)
x = tf.tanh(x)
return x
f(tf.random.uniform([10]))*[0.317767859 0.584168196 0.104753375 … 0.446179509 0.742879629 0.0446436405] [0.30748722 0.52568835 0.104371883 … 0.418753535 0.630881786 0.0446140021] …… [0.109031871 0.113768123 0.0778044686 … 0.11229489 0.114564255 0.0416750461] [0.108601853 0.113279805 0.0776478425 … 0.11182522 0.114065647 0.0416509286] \<tf.Tensor: id=1841, shape=(10,), dtype=float32, numpy= array([0.10817688, 0.11279771, 0.07749216, 0.11381008, 0.11298336, 0.10493401, 0.10051021, 0.11136142, 0.1135735 , 0.04162686], dtype=float32)> *
# 如果你很好奇，你可以检查代码签名生成。
# 但感觉就像阅读汇编语言一样。
def f(x):
while tf.reduce_sum(x) > 1:
tf.print(x)
x = tf.tanh(x)
return x
print(tf.autograph.to_code(f))from _future_ import print_function def tff(x): try: with ag.function_scope(‘f’): do_return = False retval_ = None def loop_test(x_1): with ag.function_scope(‘loop_test’): return ag.gt(ag.converted_call(‘reduce_sum’, tf, ag.ConversionOptions(recursive=True, verbose=0, strip_decorators=(ag.convert, ag.do_not_convert, ag.converted_call), force_conversion=False, optional_features=ag.Feature.ALL, internal_convert_user_code=True), (x_1,), {}), 1) def loop_body(x_1): with ag.function_scope(‘loop_body’): with ag.utils.control_dependency_on_returns(ag.converted_call(‘print’, tf, ag.ConversionOptions(recursive=True, verbose=0, strip_decorators=(ag.convert, ag.do_not_convert, ag.converted_call), force_conversion=False, optional_features=ag.Feature.ALL, internal_convert_user_code=True), (x_1,), {})): tf_1, x = ag.utils.alias_tensors(tf, x_1) x = ag.converted_call(‘tanh’, tf_1, ag.ConversionOptions(recursive=True, verbose=0, strip_decorators=(ag.convert, ag.do_not_convert, ag.converted_call), force_conversion=False, optional_features=ag.Feature.ALL, internal_convert_user_code=True), (x,), {}) return x, x, = ag.while_stmt(loop_test, loop_body, (x,), (tf, x, ag)) do_return = True retval_ = x return retval_ except: ag.rewrite_graph_construction_error(ag_source_map) tff.autograph_info__ = {}
要控制autograph，请记住它只影响Python中的基本控制流构造（if，for，while，break等），并且只有在谓词为Tensors时才会更改它们。
因此，在下面的示例中，第一个循环是静态展开的，而第二个循环是动态转换的：
@tf.function
def f(x):
for i in range(10): # Static python loop, we'll not convert it
do_stuff()
for i in tf.range(10): # depends on a tensor, we'll convert it同样，为了保证打印和断言动态发生，请使用tf.print和tf.assert：
@tf.function
def f(x):
for i in tf.range(10):
tf.print(i)
tf.Assert(i < 10, ["a"])
x += x
return x f(10)0 1 2 3 4 5 6 7 8 9 <tf.Tensor: id=1904, shape=(), dtype=int32, numpy=10240>
最后，autograph不能将任意Python代码编译成TensorFlow图。具体而言，您动态使用的数据结构仍然需要是TensorFlow数据结构。
因此，例如，在循环中累积数据的最佳方法仍然是使用tf.TensorArray：
@tf.function
def f(x):
ta = tf.TensorArray(tf.float32, size=10)
for i in tf.range(10):
x += x
ta = ta.write(i, x)
return ta.stack() f(10.0)<tf.Tensor: id=1973, shape=(10,), dtype=float32, numpy= array([ 20., 40., 80., 160., 320., 640., 1280., 2560., 5120., 10240.], dtype=float32)>
下一步
现在重新访问早期的教程并尝试使用tf.function加速代码！
