@tf.function装饰器将一个函数编译成图会发生什么？为什么它比 eager 模式下的速度更快？

Question

根据TensorFlowdocument，@tf.function将函数编译成图，"makes you get the benefits of faster execution, running on GPU or TPU, or exporting to SavedModel."

文档中的示例展示了这样的好处：

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 lstm: 0.032440788112580776 
function lstm: 0.004768412094563246

'compiled graph' 和 'function in eager mode' 有什么区别？为什么前者执行起来更快？

Answer 1

在图模式下，TensorFlow 构建一个表示您的模型的计算图，并通过会话将其转发到 C++ 运行时。这提供了分布式训练的好处并优化了过程中的计算图（通过常量折叠等）。它还简化了部署到 platform-independent 服务器的过程。

@tf.function 装饰器允许用户在急切模式下获得图形执行的好处，但请注意一些注意事项（例如，tf.function 更喜欢 TF 操作而不是 Python 基元）。

来源：https://www.tensorflow.org/guide/eager https://www.tensorflow.org/alpha/tutorials/eager/tf_function

Answer 2

在 TF 2.0 中默认使用的 Eager Execution 会立即评估操作，而无需构建图表。另一方面，图表有一定的advantages。但是，执行时间将取决于特定代码。使用 C++ 后端的代码越多，差异就越小。所有开销都来自使用 python 解释器的急切模式下的程序。例如，基本的矩阵乘法运算不会像更复杂的 DNN 应用程序那样获得更多好处。

@tf.function装饰器将一个函数编译成图会发生什么？为什么它比 eager 模式下的速度更快？

What happens when @tf.function decorator compiles a function into graph? Why is it faster than that in eager mode?

python

tensorflow

tensorflow2.0