@ tf.function装饰器将函数编译为图形时会发生什么?为什么它比渴望模式下的速度更快?

时间:2019-05-18 17:53:15

标签: python tensorflow tensorflow2.0

根据TensorFlow document@tf.function将函数编译为图形,并且“使您获得更快的执行,在GPU或TPU上运行或导出到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

“编译图”和“急切模式下的功能”之间有什么区别?为什么前者执行时会更快?

2 个答案:

答案 0 :(得分:1)

在图模式下,TensorFlow会构建一个表示您的模型的计算图,并通过会话将其转发到C ++运行时。这样可提供分布式培训收益,并在此过程中(通过不断折叠等)优化您的计算图。它还简化了部署到与平台无关的服务器的过程。

@ tf.function装饰器允许用户在热切的模式下获得图执行的好处,但请注意,有一些注意事项(例如,tf.function比Python原语更喜欢TF ops)。

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

答案 1 :(得分:1)

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